uniugm/admin/cmd/tools/generator/errcode/gen.go

package errcode

import (
	"flag"
	"fmt"
	"go/ast"
	"go/token"
	"golang.org/x/tools/go/packages"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

type File struct {
	pkg         *Package
	fileName    string
	rawAstFile  *ast.File
	content     string
	importList  map[string]struct{}
	errcodeList *ErrCodeList
}

type Package struct {
	Name          string
	rawAstPkg     *packages.Package
	rawAstFileSet *token.FileSet
	files         []*File
	ImportList    []string
}

func Generate() {
	flag.Parse()

	curRootDir, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("================================================")
	log.Printf("args: %v", os.Args[1:])
	log.Printf("pwd: %v", curRootDir)

	cfg := &packages.Config{
		Mode:       packages.NeedName | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedSyntax | packages.NeedFiles,
		Tests:      true,
		BuildFlags: []string{fmt.Sprintf("-tags=%s", strings.Join([]string{}, " "))},
		Logf: func(format string, args ...interface{}) {
			//log.Printf(format, args...)
		},
	}

	rawPkgs, err := packages.Load(cfg, ".")
	if err != nil {
		log.Fatal(err)
	}
	if len(rawPkgs) == 0 {
		log.Fatalf("当前包居然读取不到ast信息，尝试项目根目录执行go mod tidy，解决报错")
	}

	pkgs := make([]*Package, 0, len(rawPkgs))
	firstPkgPath := ""
	for i, pkg := range rawPkgs {
		if i == 0 {
			firstPkgPath = pkg.PkgPath
		} else {
			if strings.Contains(pkg.PkgPath, firstPkgPath) {
				continue
			}
		}
		p := &Package{
			Name:          pkg.Name,
			rawAstPkg:     pkg,
			rawAstFileSet: pkg.Fset,
			files:         make([]*File, 0, len(pkg.Syntax)),
		}

		for j, file := range pkg.Syntax {
			goFile := pkg.GoFiles[j]

			if strings.HasSuffix(goFile, ".gen.go") {
				continue
			}

			content, err := os.ReadFile(goFile)
			if err != nil {
				panic(err)
			}

			p.files = append(p.files, &File{
				pkg:        p,
				fileName:   filepath.Base(goFile),
				rawAstFile: file,
				content:    string(content),
				importList: make(map[string]struct{}),
			})
		}
		pkgs = append(pkgs, p)
	}

	if len(pkgs) > 1 {
		s := ""
		for _, v := range pkgs {
			s += v.rawAstPkg.Name + ","
			log.Printf("pak:%v, gofiles:%v", v.rawAstPkg.PkgPath, v.rawAstPkg.GoFiles)
		}
		log.Fatalf("居然发现%v个包(%v)，联系@李坤支持:", len(pkgs), s)
	}

	for _, pkg := range pkgs {
		for _, file := range pkg.files {
			file.errcodeList = &ErrCodeList{
				List: make([]*ErrCode, 0),
			}
			ast.Inspect(file.rawAstFile, file.genDecl)

			//log.Printf("after gen, file %v list len:%v", file.fileName, len(file.errcodeList.List))

			if len(file.errcodeList.List) > 0 {
				fileName := "code_content.gen.go"
				outputFile := filepath.Join(curRootDir, fileName)
				err := renderFile(file.errcodeList, tpl, outputFile)
				if err != nil {
					log.Printf("render file %v error:%v", fileName, err)
					os.Exit(1)
				}
			}
		}
	}

}

// genDecl processes one declaration clause.
func (f *File) genDecl(node ast.Node) bool {
	//log.Printf("node:%+v", node)
	if f.pkg.Name != "errcode" {
		return true
	}
	if f.fileName != "code.go" {
		return true
	}

	if node == nil {
		return true
	}

	fileSet := f.pkg.rawAstFileSet
	fpos := fileSet.File(node.Pos())

	fileNode, ok := node.(*ast.File)
	if !ok {
		return true
	}

	for _, v := range fileNode.Decls {
		d := v.(*ast.GenDecl)

		if d.Tok.String() != "const" {
			continue
		}

		// 静态错误码
		for _, spec := range d.Specs {
			specValue := spec.(*ast.ValueSpec)

			specContent := f.content[fpos.Offset(specValue.Pos()):fpos.Offset(specValue.End())]
			specContent = strings.ReplaceAll(specContent, " ", "")
			values := strings.Split(specContent, "=")
			if len(values) != 2 {
				log.Printf("错误码：%v 按等号分隔字符串之后不满足var=value格式", specContent)
				os.Exit(1)
			}

			commentContent := specValue.Comment.Text()
			commentContent = strings.ReplaceAll(commentContent, "\n", "")
			if commentContent == "" {
				log.Printf("错误码:%v 没有定义文本内容，格式为：ErrCode = 123 // <文本内容>", specContent)
				os.Exit(1)
			}

			code, _ := strconv.Atoi(values[1])
			f.errcodeList.List = append(f.errcodeList.List, &ErrCode{
				Name:    values[0],
				Value:   code,
				Comment: commentContent,
			})

			//log.Printf("错误码：%v 值：%v --> 注释：%v\n", values[0], values[1], commentContent)
		}

	}

	return true
}