github.com/cloudwego/kitex@v0.9.0/tool/internal_pkg/generator/template.go

// Copyright 2022 CloudWeGo Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package generator

import (
	"encoding/json"
	"fmt"
	"io/ioutil"

	"gopkg.in/yaml.v3"
)

// APIExtension contains segments to extend an API template.
type APIExtension struct {
	// ImportPaths contains a list of import path that the file should add to the import list.
	// The paths must be registered with the TemplateExtension's Dependencies fields.
	ImportPaths []string `json:"import_paths,omitempty" yaml:"import_paths,omitempty"`

	// Code snippets to be inserted in the NewX function before assembling options.
	// It must be a template definition with a name identical to the one kitex uses.
	ExtendOption string `json:"extend_option,omitempty" yaml:"extend_option,omitempty"`

	// Code snippets to be appended to the file.
	// It must be a template definition with a name identical to the one kitex uses.
	ExtendFile string `json:"extend_file,omitempty" yaml:"extend_file,omitempty"`
}

// TemplateExtension extends templates that generates files in *service packages.
type TemplateExtension struct {
	// FeatureNames registers some names to the scope for the code generating phrase, where templates can use the `HasFeature` function to query.
	FeatureNames []string `json:"feature_names,omitempty" yaml:"feature_names,omitempty"`

	// EnableFeatures marks on which features that `HasFeature` queries should return true.
	EnableFeatures []string `json:"enable_features,omitempty" yaml:"enable_features,omitempty"`

	// Dependencies is a mapping from import path to package names/aliases.
	Dependencies map[string]string `json:"dependencies,omitempty" yaml:"dependencies,omitempty"`

	// Extension for client.go .
	ExtendClient *APIExtension `json:"extend_client,omitempty" yaml:"extend_client,omitempty"`

	// Extension for server.go .
	ExtendServer *APIExtension `json:"extend_server,omitempty" yaml:"extend_server,omitempty"`

	// Extension for invoker.go .
	ExtendInvoker *APIExtension `json:"extend_invoker,omitempty" yaml:"extend_invoker,omitempty"`
}

// FromJSONFile unmarshals a TemplateExtension with JSON format from the given file.
func (p *TemplateExtension) FromJSONFile(filename string) error {
	if p == nil {
		return nil
	}
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, p)
}

// ToJSONFile marshals a TemplateExtension to the given file in JSON format.
func (p *TemplateExtension) ToJSONFile(filename string) error {
	data, err := json.Marshal(p)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filename, data, 0o644)
}

// FromYAMLFile unmarshals a TemplateExtension with YAML format from the given file.
func (p *TemplateExtension) FromYAMLFile(filename string) error {
	if p == nil {
		return nil
	}
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}
	return yaml.Unmarshal(data, p)
}

func (p *TemplateExtension) ToYAMLFile(filename string) error {
	data, err := yaml.Marshal(p)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filename, data, 0o644)
}

func (p *TemplateExtension) Merge(other *TemplateExtension) {
	if other == nil {
		return
	}
	p.FeatureNames = append(p.FeatureNames, other.FeatureNames...)
	p.EnableFeatures = append(p.EnableFeatures, other.EnableFeatures...)

	if other.Dependencies != nil {
		if p.Dependencies == nil {
			p.Dependencies = other.Dependencies
		} else {
			for k, v := range other.Dependencies {
				p.Dependencies[k] = v
			}
		}
	}

	if other.ExtendClient != nil {
		if p.ExtendClient == nil {
			p.ExtendClient = other.ExtendClient
		} else {
			p.ExtendClient.Merge(other.ExtendClient)
		}
	}
	if other.ExtendServer != nil {
		if p.ExtendServer == nil {
			p.ExtendServer = other.ExtendServer
		} else {
			p.ExtendServer.Merge(other.ExtendServer)
		}
	}
	if other.ExtendInvoker != nil {
		if p.ExtendInvoker == nil {
			p.ExtendInvoker = other.ExtendInvoker
		} else {
			p.ExtendInvoker.Merge(other.ExtendInvoker)
		}
	}
}

func (a *APIExtension) Merge(other *APIExtension) {
	if other == nil {
		return
	}
	a.ImportPaths = append(a.ImportPaths, other.ImportPaths...)
	a.ExtendOption = fmt.Sprintf("%v\n%v", a.ExtendOption, other.ExtendOption)
	a.ExtendFile = fmt.Sprintf("%v\n%v", a.ExtendFile, other.ExtendFile)
}