generator.go

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"io/fs"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
)

type Generator struct {
	buf    bytes.Buffer
	files  []string
	output string
}

func (g *Generator) generate(dir string) {
	// 获取所有目录下的文件,放入g.files中
	g.parseDir(dir)
	// 获取所有文件中所有的结构体以及字段(目前来看字段没啥用以后,
	// 但是生成关联的时候必然有用
	files, err := g.parseFile()
	if err != nil {
		panic(err)
	}
	for _, file := range files {
		// 生成需要引入的包以及固定的注释
		g.buildPrepare(pkgPath(file.dir))
		for _, s := range file.structs {
			// 生成dao层create
			g.buildCreate(file.pkg, s)
			// 生成dao层list
			g.buildList(file.pkg, s)
			// 生成dao层query
			g.buildQuery(file.pkg, s)
			// 生成dao层scope
			g.buildScope(file.pkg, s)
			// 生成dao层update
			g.buildUpdate(file.pkg, s)
			// 生成dao层delete
			g.buildDelete(file.pkg, s)
		}
		// 格式化代码
		src := g.format()
		// 确认生成文件
		outputFile := path.Join(g.output, file.name)
		// 输出到指定文件
		err := os.WriteFile(outputFile, src, 0644)
		if err != nil {
			panic(err)
		}
		// 重置
		g.reset()
	}
}

// ParseDir 获取该目录下的所有文件
func (g *Generator) parseDir(dir string) {
	if strings.HasSuffix(dir, ".go") {
		return
	}
	err := filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
		if !d.IsDir() {
			g.files = append(g.files, path)
		}
		return nil
	})
	if err != nil {
		panic(err)
	}
}

// ParseFile 解析文件获得所有文件里的结构体以及包名称以及生成文件名称
func (g *Generator) parseFile() ([]*File, error) {
	if len(g.files) == 0 {
		return nil, nil
	}
	files := make([]*File, 0)
	for _, f := range g.files {
		// 解析文件
		file := &File{structs: make([]*Struct, 0)}
		// 获取文件的路径以及名称
		file.dir, file.name = filename(f)
		// 解析package名
		fset := token.NewFileSet()
		parse, err := parser.ParseFile(fset, f, nil, parser.PackageClauseOnly)
		if err != nil {
			return nil, err
		}
		file.pkg = parse.Name.Name

		parse, err = parser.ParseFile(fset, f, nil, 0)
		if err != nil {
			return nil, err
		}
		// 遍历节点
		ast.Inspect(parse, func(node ast.Node) bool {
			switch t := node.(type) {
			// 如果是声明type
			case *ast.TypeSpec:
				// 确认这个type的声明是否是struct类型
				structType, ok := t.Type.(*ast.StructType)
				// 获取所有的字段名称、类型、标签
				if ok {
					stru := &Struct{}
					stru.name = t.Name.Name
					fields := make([]*Field, 0)
					for _, v := range structType.Fields.List {
						if v.Tag == nil {
							continue
						}
						tags := extractTag(v.Tag.Value)
						typIdent := v.Type.(*ast.Ident)
						var typ string
						if typIdent.Obj != nil {
							typ = file.pkg + "." + typIdent.Name
						} else {
							typ = typIdent.Name
						}
						if tags.Gorm(v.Names[0].Name) {
							fields = append(fields, &Field{
								name:   v.Names[0].Name,
								goType: typ,
								tags:   tags,
							})
						}
					}
					stru.fields = fields
					file.structs = append(file.structs, stru)
				}
			}
			return true
		})
		files = append(files, file)
	}
	return files, nil
}

// buildPrepare 每次生成都需要
func (g *Generator) buildPrepare(pkgPath string) {
	pkgName := g.output[strings.LastIndex(g.output, "/")+1:]
	g.Printf("// Code generated by \"gorm-gen\"; DO NOT EDIT.\n")
	g.Printf("\n")
	g.Printf("package %s\n", pkgName)
	g.Printf("import (\n")
	g.Printf("\"context\"\n")
	g.Printf("\"%s\"\n", pkgPath)
	g.Printf("\"gorm.io/gorm\"\n")
	g.Printf(")")
	g.Printf("\n")
}

// format 使用gofmt格式化go代码
func (g *Generator) format() []byte {
	src, err := format.Source(g.buf.Bytes())
	if err != nil {
		log.Printf("warning: internal error: invalid Go generated: %s", err)
		log.Printf("warning: compile the package to analyze the error")
		return g.buf.Bytes()
	}
	return src
}

// Printf 格式输出
func (g *Generator) Printf(format string, args ...interface{}) {
	_, _ = fmt.Fprintf(&g.buf, format, args...)
}

// reset 重置buffer
func (g *Generator) reset() {
	g.buf.Reset()
}