provider.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"

	"github.com/fatih/structtag"
	"golang.org/x/tools/imports"
)

const (
	AccessRead    = "r"
	AccessWrite   = "w"
	AccessTagName = "access"

	ModeTypeAll    = 0
	ModeTypeGetter = 1
	ModeTypeSetter = 2
)

var (
	fileName = flag.String("file", "", "a parsed filename must be set if you use command line")
	modeType = flag.Int("mode", 0, "when mode=0 or not set, gen Getter and Setter!\n"+
		"when mode=1, only gen Getter method!\n"+
		"when mode=2, only gen Setter method!\n"+
		"when mode is other num, not support!")
)

func main() {
	log.SetFlags(0) // 设置日志的抬头信息
	log.SetPrefix("gentools-accessor: ")
	flag.Parse() //把用户传递的命令行参数解析为对应变量的值

	var mode int
	if modeType != nil {
		if *modeType < 0 || *modeType > 2 {
			log.Fatalln("\n\tThe -mode parameter setting is out of range, please execute the command to check the meaning of the parameter: \n\t\t>>>\tgentools-accessor -h")
			return
		} else {
			mode = *modeType
		}
	}

	var inputName string
	if len(*fileName) > 0 {
		inputName = *fileName
	} else {
		//尝试获取当前被执行的文件
		inputName = os.Getenv("GOFILE")
	}

	if len(inputName) == 0 {
		log.Fatalln("Please enter a correct file path to be parsed！")
		return
	}
	log.Println("\n\tThe currently executed fileName is: ", inputName)

	g := Generator{
		buf: bytes.NewBufferString(""),
	}
	g.generate(inputName, mode)
	src := g.formatCode()
	outputName := strings.TrimSuffix(inputName, ".go") + ".accessor.go"
	err := ioutil.WriteFile(outputName, src, 0644) // ignore_security_alert
	if err != nil {
		log.Fatalf("writing output: %s\n", err)
		return
	}
	log.Printf("\n\t🎉🎉🎉 Congratulation！\n\tAutomatic code generation finished!\n\tOutput_name: %s\n", outputName)
}

type StructFieldInfo struct {
	Name   string   // 字段名
	Type   string   // 类型名
	Access []string // tag对应的value，即为r，w
}

type StructFieldInfoArr = []StructFieldInfo

type Generator struct {
	buf *bytes.Buffer // Accumulated output.
}

// FormatCode sets the options of the imports pkg and then applies the Process method
// which by default removes all of the imports not used and formats the remaining docs,
// imports and code like `gofmt`. It will e.g. remove paranthesis around a unnamed single return type
func (g *Generator) formatCode() []byte {
	opts := &imports.Options{
		TabIndent: true,
		TabWidth:  2,
		Fragment:  true,
		Comments:  true,
	}
	src := (g.buf).Bytes()
	fmtCode, err := imports.Process("", src, opts)
	if err != nil {
		log.Fatalln("格式化代码失败")
		return src
	}
	return fmtCode
}

func (g *Generator) myPrintf(format string, args ...interface{}) {
	_, _ = fmt.Fprintf(g.buf, format, args...)
}

func (g *Generator) generate(fileName string, modeVal int) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, fileName, nil, parser.ParseComments)
	if err != nil || f == nil {
		log.Fatalln("parse file err")
		return
	}
	structInfo, err := parseAllStructInSingleFile(f, fset, AccessTagName, modeVal)
	if err != nil {
		log.Fatalln("parse structs in file err: ", err)
		return
	}
	g.myPrintf("// Code generated by \"gentools-accessor\"; DO NOT EDIT.\n")
	g.myPrintf("\n")
	g.myPrintf("package %s\n", f.Name)
	g.myPrintf("\n")
	// 将源文件导入的包名写入
	if len(f.Imports) > 0 {
		g.myPrintf("%s\n", genImports(f.Imports))
	}
	for stName, info := range structInfo {
		// 这里可以维护一个map，对仅仅是大小写字母区别的做标记
		visMap := make(map[string]int)
		for index, field := range info {
			lowerName := strings.ToLower(field.Name)
			if d, exist := visMap[lowerName]; exist {
				log.Fatalf("\n\t结构体 %s 中：有两个含义相近的命名：%s 和 %s，请改正！终止生成代码～\n", stName, info[d].Name, field.Name)
				return
			}
			visMap[lowerName] = index
			for _, access := range field.Access {
				switch access {
				case AccessWrite:
					g.myPrintf("%s\n", genSetter(stName, field.Name, field.Type))
				case AccessRead:
					g.myPrintf("%s\n", genGetter(stName, field.Name, field.Type))
				}
			}
		}
	}
}

// firstUpper 字符串首字母大写
func firstUpper(s string) string {
	if s == "" {
		return ""
	}
	return strings.ToUpper(s[:1]) + s[1:]
}

func parseAllStructInSingleFile(file *ast.File, fileSet *token.FileSet, tagName string, modeVal int) (structMap map[string]StructFieldInfoArr, err error) {
	structMap = make(map[string]StructFieldInfoArr)

	collectStructs := func(x ast.Node) bool {
		ts, ok := x.(*ast.TypeSpec)
		if !ok || ts.Type == nil {
			return true
		}

		// 获取结构体名称
		structName := ts.Name.Name

		s, ok := ts.Type.(*ast.StructType)
		if !ok {
			return true
		}
		fileInfos := make([]StructFieldInfo, 0)
		for _, field := range s.Fields.List {
			if len(field.Names) == 0 { // Notice: 跳过组合struct的字段
				continue
			}
			name := field.Names[0].Name // 字段名称
			info := StructFieldInfo{Name: name}
			var typeNameBuf bytes.Buffer
			err := printer.Fprint(&typeNameBuf, fileSet, field.Type)
			if err != nil {
				log.Println("获取字段类型失败:", err)
				return true
			}
			info.Type = typeNameBuf.String()
			if field.Tag != nil { // 有tag
				tag := field.Tag.Value
				tag = strings.Trim(tag, "`")
				tags, err := structtag.Parse(tag)
				if err != nil {
					return true
				}
				access, err := tags.Get(tagName)
				if err == nil {
					access.Options = append(access.Options, access.Name)
					for i, v := range access.Options {
						if v == AccessRead || v == AccessWrite {
							continue
						}
						// 剔除除 r,w 之外的 tag-value
						access.Options = append(access.Options[:i], access.Options[i+1:]...)
					}
				}
				info.Access = access.Options
			} else {
				if modeVal == ModeTypeAll { // 模式为0，封装可读可读方法
					info.Access = []string{AccessRead, AccessWrite}
				} else if modeVal == ModeTypeGetter { // 模式为1，只封装可读方法
					info.Access = []string{AccessRead}
				} else if modeVal == ModeTypeSetter { // 模式为2，只封装可写方法
					info.Access = []string{AccessWrite}
				}
			}
			fileInfos = append(fileInfos, info)
		}
		structMap[structName] = fileInfos
		return false
	}

	ast.Inspect(file, collectStructs)

	return structMap, nil
}

func genImports(importsArr []*ast.ImportSpec) string {
	packageNameArr := make([]string, 0)
	for _, s := range importsArr {
		if s.Name != nil {
			packageNameArr = append(packageNameArr, fmt.Sprintf("%s %s\n", s.Name.Name, s.Path.Value))
		} else {
			packageNameArr = append(packageNameArr, s.Path.Value)
		}
	}
	tpl := `import (
		{{range .}}
			{{- .}}
		{{end -}}
	)`
	t := template.New("imports")
	t = template.Must(t.Parse(tpl))
	res := bytes.NewBufferString("")
	_ = t.Execute(res, packageNameArr)
	return res.String()
}

func genSetter(structName, fieldName, typeName string) string {
	tpl := `func ({{.Receiver}} *{{.Struct}}) Set{{.MethodName}}(param {{.Type}}) {
	{{.Receiver}}.{{.Field}} = param
}`
	t := template.New("setter")
	t = template.Must(t.Parse(tpl))
	res := bytes.NewBufferString("")
	_ = t.Execute(res, map[string]string{
		"Receiver":   strings.ToLower(structName[0:1]),
		"Struct":     structName,
		"Field":      fieldName,
		"Type":       typeName,
		"MethodName": firstUpper(fieldName),
	})
	return res.String()
}

func genGetter(structName, fieldName, typeName string) string {
	tpl := `func ({{.Receiver}} *{{.Struct}}) Get{{.MethodName}}() (v0 {{.Type}}) {
	if {{.Receiver}} == nil {
		return
	}
	return {{.Receiver}}.{{.Field}}
}`
	t := template.New("getter")
	t = template.Must(t.Parse(tpl))
	res := bytes.NewBufferString("")
	_ = t.Execute(res, map[string]string{
		"Receiver":   strings.ToLower(structName[0:1]),
		"Struct":     structName,
		"Field":      fieldName,
		"Type":       typeName,
		"MethodName": firstUpper(fieldName),
	})
	return res.String()
}