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

// Copyright 2021 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 (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"

	"gopkg.in/yaml.v3"

	"github.com/cloudwego/kitex/tool/internal_pkg/log"
	"github.com/cloudwego/kitex/tool/internal_pkg/util"
)

var DefaultDelimiters = [2]string{"{{", "}}"}

type updateType string

const (
	skip              updateType = "skip"
	cover             updateType = "cover"
	incrementalUpdate updateType = "append"
)

type Update struct {
	// update type: skip / cover / append. Default is skip.
	// If `LoopMethod` is true, only Type field is effect and no append behavior.
	Type string `yaml:"type,omitempty"`
	// Match key in append type. If the rendered key exists in the file, the method will be skipped.
	Key string `yaml:"key,omitempty"`
	// Append template. Use it to render append content.
	AppendTpl string `yaml:"append_tpl,omitempty"`
	// Append import template. Use it to render import content to append.
	ImportTpl []string `yaml:"import_tpl,omitempty"`
}

type Template struct {
	// The generated path and its filename. For example: biz/test.go
	// will generate test.go in biz directory.
	Path string `yaml:"path,omitempty"`
	// Render template content, currently only supports go template syntax
	Body string `yaml:"body,omitempty"`
	// define update behavior
	UpdateBehavior *Update `yaml:"update_behavior,omitempty"`
	// If set this field, kitex will generate file by cycle. For example:
	// test_a/test_b/{{ .Name}}_test.go
	LoopMethod bool `yaml:"loop_method,omitempty"`
	// If both set this field and combine-service, kitex will generate service by cycle.
	LoopService bool `yaml:"loop_service,omitempty"`
}

type customGenerator struct {
	fs       []*File
	pkg      *PackageInfo
	basePath string
}

func NewCustomGenerator(pkg *PackageInfo, basePath string) *customGenerator {
	return &customGenerator{
		pkg:      pkg,
		basePath: basePath,
	}
}

func (c *customGenerator) loopGenerate(tpl *Template) error {
	tmp := c.pkg.Methods
	defer func() {
		c.pkg.Methods = tmp
	}()
	m := c.pkg.AllMethods()
	for _, method := range m {
		c.pkg.Methods = []*MethodInfo{method}
		pathTask := &Task{
			Text: tpl.Path,
		}
		renderPath, err := pathTask.RenderString(c.pkg)
		if err != nil {
			return err
		}
		filePath := filepath.Join(c.basePath, renderPath)
		// update
		if util.Exists(filePath) && updateType(tpl.UpdateBehavior.Type) == skip {
			continue
		}
		task := &Task{
			Name: path.Base(renderPath),
			Path: filePath,
			Text: tpl.Body,
		}
		f, err := task.Render(c.pkg)
		if err != nil {
			return err
		}
		c.fs = append(c.fs, f)
	}
	return nil
}

func (c *customGenerator) commonGenerate(tpl *Template) error {
	// Use all services including base service.
	tmp := c.pkg.Methods
	defer func() {
		c.pkg.Methods = tmp
	}()
	c.pkg.Methods = c.pkg.AllMethods()

	pathTask := &Task{
		Text: tpl.Path,
	}
	renderPath, err := pathTask.RenderString(c.pkg)
	if err != nil {
		return err
	}
	filePath := filepath.Join(c.basePath, renderPath)
	update := util.Exists(filePath)
	if update && updateType(tpl.UpdateBehavior.Type) == skip {
		log.Infof("skip generate file %s", tpl.Path)
		return nil
	}
	var f *File
	if update && updateType(tpl.UpdateBehavior.Type) == incrementalUpdate {
		cc := &commonCompleter{
			path:   filePath,
			pkg:    c.pkg,
			update: tpl.UpdateBehavior,
		}
		f, err = cc.Complete()
		if err != nil {
			return err
		}
	} else {
		// just create dir
		if tpl.Path[len(tpl.Path)-1] == '/' {
			os.MkdirAll(filePath, 0o755)
			return nil
		}

		task := &Task{
			Name: path.Base(tpl.Path),
			Path: filePath,
			Text: tpl.Body,
		}

		f, err = task.Render(c.pkg)
		if err != nil {
			return err
		}
	}

	c.fs = append(c.fs, f)
	return nil
}

func (g *generator) GenerateCustomPackage(pkg *PackageInfo) (fs []*File, err error) {
	g.updatePackageInfo(pkg)

	g.setImports(HandlerFileName, pkg)
	t, err := readTemplates(g.TemplateDir)
	if err != nil {
		return nil, err
	}
	for _, tpl := range t {
		if tpl.LoopService && g.CombineService {
			svrInfo, cs := pkg.ServiceInfo, pkg.CombineServices

			for i := range cs {
				pkg.ServiceInfo = cs[i]
				f, err := renderFile(pkg, g.OutputPath, tpl)
				if err != nil {
					return nil, err
				}
				fs = append(fs, f...)
			}
			pkg.ServiceInfo, pkg.CombineServices = svrInfo, cs
		} else {
			f, err := renderFile(pkg, g.OutputPath, tpl)
			if err != nil {
				return nil, err
			}
			fs = append(fs, f...)
		}
	}
	return fs, nil
}

func renderFile(pkg *PackageInfo, outputPath string, tpl *Template) (fs []*File, err error) {
	cg := NewCustomGenerator(pkg, outputPath)
	// special handling Methods field
	if tpl.LoopMethod {
		err = cg.loopGenerate(tpl)
	} else {
		err = cg.commonGenerate(tpl)
	}
	if err == errNoNewMethod {
		err = nil
	}
	return cg.fs, err
}

func readTemplates(dir string) ([]*Template, error) {
	files, _ := ioutil.ReadDir(dir)
	var ts []*Template
	for _, f := range files {
		// filter dir and non-yaml files
		if f.Name() != ExtensionFilename && !f.IsDir() && (strings.HasSuffix(f.Name(), "yaml") || strings.HasSuffix(f.Name(), "yml")) {
			p := filepath.Join(dir, f.Name())
			tplData, err := ioutil.ReadFile(p)
			if err != nil {
				return nil, fmt.Errorf("read layout config from  %s failed, err: %v", p, err.Error())
			}
			t := &Template{
				UpdateBehavior: &Update{Type: string(skip)},
			}
			if err = yaml.Unmarshal(tplData, t); err != nil {
				return nil, fmt.Errorf("%s: unmarshal layout config failed, err: %s", f.Name(), err.Error())
			}
			ts = append(ts, t)
		}
	}

	return ts, nil
}