#11 code generator

cmd/gormgen/README.MD

@@ -0,0 +1,6 @@
+## 执行命令
+1. 定义生成的表，设置 config 中 cmd.genTables，可以自定义设置多张表，为空表示生成库中所有的表，如果设置多个表可用','分割；
+1. 在根目录下执行脚本文件：`./scripts/gormgen.sh`；
+
+## 参考
+- https://github.com/MohamedBassem/gormgen

cmd/gormgen/main.go

@@ -0,0 +1,37 @@
+package main
+
+import (
+	"flag"
+	"log"
+	"os"
+	"strings"
+
+	"github.com/xinliangnote/go-gin-api/cmd/gormgen/pkg"
+)
+
+var (
+	input   string
+	structs []string
+)
+
+func init() {
+	flagStructs := flag.String("structs", "", "[Required] The name of schema structs to generate structs for, comma seperated\n")
+	flagInput := flag.String("input", "", "[Required] The name of the input file dir\n")
+	flag.Parse()
+
+	if *flagStructs == "" || *flagInput == "" {
+		flag.Usage()
+		os.Exit(1)
+	}
+
+	structs = strings.Split(*flagStructs, ",")
+	input = *flagInput
+}
+
+func main() {
+	gen := pkg.NewGenerator(input)
+	p := pkg.NewParser(input)
+	if err := gen.ParserAST(p, structs).Generate().Format().Flush(); err != nil {
+		log.Fatalln(err)
+	}
+}

cmd/gormgen/pkg/generator.go

@@ -0,0 +1,124 @@
+package pkg
+
+import (
+	"bytes"
+	"errors"
+	"go/format"
+	"io/ioutil"
+	"log"
+	"strings"
+
+	"github.com/jinzhu/gorm"
+)
+
+// fieldConfig
+type fieldConfig struct {
+	FieldName  string
+	ColumnName string
+	FieldType  string
+	HumpName   string
+}
+
+// structConfig
+type structConfig struct {
+	config
+	StructName   string
+	OnlyFields   []fieldConfig
+	OptionFields []fieldConfig
+}
+
+type ImportPkg struct {
+	Pkg string
+}
+
+type structHelpers struct {
+	Titelize func(string) string
+}
+
+type config struct {
+	PkgName          string
+	Helpers          structHelpers
+	QueryBuilderName string
+}
+
+// The Generator is the one responsible for generating the code, adding the imports, formating, and writing it to the file.
+type Generator struct {
+	buf           map[string]*bytes.Buffer
+	inputFile     string
+	config        config
+	structConfigs []structConfig
+}
+
+// NewGenerator function creates an instance of the generator given the name of the output file as an argument.
+func NewGenerator(outputFile string) *Generator {
+	return &Generator{
+		buf:       map[string]*bytes.Buffer{},
+		inputFile: outputFile,
+	}
+}
+
+// ParserAST parse by go file
+func (g *Generator) ParserAST(p *Parser, structs []string) (ret *Generator) {
+	for _, v := range structs {
+		g.buf[gorm.ToDBName(v)] = new(bytes.Buffer)
+	}
+	g.structConfigs = p.Parse()
+	g.config.PkgName = p.pkg.Name
+	g.config.Helpers = structHelpers{
+		Titelize: strings.Title,
+	}
+	g.config.QueryBuilderName = SQLColumnToHumpStyle(p.pkg.Name) + "QueryBuilder"
+	return g
+}
+
+func (g *Generator) checkConfig() (err error) {
+	if len(g.config.PkgName) == 0 {
+		err = errors.New("package name dose'n set")
+		return
+	}
+	for i := 0; i < len(g.structConfigs); i++ {
+		g.structConfigs[i].config = g.config
+	}
+	return
+}
+
+// Generate executes the template and store it in an internal buffer.
+func (g *Generator) Generate() *Generator {
+	if err := g.checkConfig(); err != nil {
+		panic(err)
+	}
+
+	for _, v := range g.structConfigs {
+		if _, ok := g.buf[gorm.ToDBName(v.StructName)]; !ok {
+			continue
+		}
+		if err := outputTemplate.Execute(g.buf[gorm.ToDBName(v.StructName)], v); err != nil {
+			panic(err)
+		}
+	}
+
+	return g
+}
+
+// Format function formats the output of the generation.
+func (g *Generator) Format() *Generator {
+	for k := range g.buf {
+		formattedOutput, err := format.Source(g.buf[k].Bytes())
+		if err != nil {
+			panic(err)
+		}
+		g.buf[k] = bytes.NewBuffer(formattedOutput)
+	}
+	return g
+}
+
+// Flush function writes the output to the output file.
+func (g *Generator) Flush() error {
+	for k := range g.buf {
+		filename := g.inputFile + "/gen_" + strings.ToLower(k) + ".go"
+		if err := ioutil.WriteFile(filename, g.buf[k].Bytes(), 0777); err != nil {
+			log.Fatalln(err)
+		}
+	}
+	return nil
+}

cmd/gormgen/pkg/parser.go

@@ -0,0 +1,120 @@
+package pkg
+
+import (
+	"go/ast"
+	"go/build"
+	"go/parser"
+	"go/token"
+	"log"
+	"strings"
+
+	"github.com/jinzhu/gorm"
+)
+
+// The Parser is used to parse a directory and expose information about the structs defined in the files of this directory.
+type Parser struct {
+	dir         string
+	pkg         *build.Package
+	parsedFiles []*ast.File
+}
+
+// NewParser create a new parser instance.
+func NewParser(dir string) *Parser {
+	return &Parser{
+		dir: dir,
+	}
+}
+
+// getPackage parse dir get go file and package
+func (p *Parser) getPackage() {
+	pkg, err := build.Default.ImportDir(p.dir, build.ImportComment)
+	if err != nil {
+		log.Fatalf("cannot process directory %s: %s", p.dir, err)
+	}
+	p.pkg = pkg
+
+}
+
+// parseGoFiles parse go file
+func (p *Parser) parseGoFiles() {
+	var parsedFiles []*ast.File
+	fs := token.NewFileSet()
+	for _, file := range p.pkg.GoFiles {
+		file = p.dir + "/" + file
+		parsedFile, err := parser.ParseFile(fs, file, nil, 0)
+		if err != nil {
+			log.Fatalf("parsing package: %s: %s\n", file, err)
+		}
+		parsedFiles = append(parsedFiles, parsedFile)
+	}
+	p.parsedFiles = parsedFiles
+}
+
+// parseTypes parse type of struct
+func (p *Parser) parseTypes(file *ast.File) (ret []structConfig) {
+	ast.Inspect(file, func(n ast.Node) bool {
+		decl, ok := n.(*ast.GenDecl)
+		if !ok || decl.Tok != token.TYPE {
+			return true
+		}
+
+		for _, spec := range decl.Specs {
+			var (
+				data structConfig
+			)
+			typeSpec, _ok := spec.(*ast.TypeSpec)
+			if !_ok {
+				continue
+			}
+			// We only care about struct declaration (for now)
+			var structType *ast.StructType
+			if structType, ok = typeSpec.Type.(*ast.StructType); !ok {
+				continue
+			}
+
+			data.StructName = typeSpec.Name.Name
+			for _, v := range structType.Fields.List {
+				var (
+					optionField fieldConfig
+				)
+
+				// type is ident, get onlyField type
+				if t, _ok := v.Type.(*ast.Ident); _ok {
+					optionField.FieldType = t.String()
+				} else {
+					if v.Tag != nil {
+						if strings.Contains(v.Tag.Value, "gorm") && strings.Contains(v.Tag.Value, "time") {
+							optionField.FieldType = "time.Time"
+						}
+					}
+				}
+
+				// get file name
+				if len(v.Names) > 0 {
+					optionField.FieldName = v.Names[0].String()
+					optionField.ColumnName = gorm.ToDBName(optionField.FieldName)
+					optionField.HumpName = SQLColumnToHumpStyle(optionField.ColumnName)
+				}
+
+				data.OptionFields = append(data.OptionFields, optionField)
+			}
+
+			ret = append(ret, data)
+		}
+		return true
+	})
+	return
+}
+
+// Parse should be called before any type querying for the parser. It takes the directory to be parsed and extracts all the structs defined in this directory.
+func (p *Parser) Parse() (ret []structConfig) {
+	var (
+		data []structConfig
+	)
+	p.getPackage()
+	p.parseGoFiles()
+	for _, f := range p.parsedFiles {
+		data = append(data, p.parseTypes(f)...)
+	}
+	return data
+}