1.8.4 元编程与代码生成 #

// example.go
package main

//go:generate stringer -type=Status
type Status int

const (
    Pending Status = iota
    Running
    Completed
    Failed
)

//go:generate go run generator.go
type User struct {
    ID   int    `json:"id"`
    Name string `json:"name"`
    Age  int    `json:"age"`
}

func main() {
    status := Running
    println(status.String()) // 这个方法将由 stringer 工具生成
}

// generator.go
package main

import (
    "fmt"
    "go/ast"
    "go/parser"
    "go/token"
    "os"
    "strings"
    "text/template"
)

// StructInfo 结构体信息
type StructInfo struct {
    Name   string
    Fields []FieldInfo
}

// FieldInfo 字段信息
type FieldInfo struct {
    Name string
    Type string
    Tag  string
}

// 生成器模板
const builderTemplate = `// Code generated by generator.go; DO NOT EDIT.

package main

{{range .}}
// {{.Name}}Builder 为 {{.Name}} 提供构建器模式
type {{.Name}}Builder struct {
    instance *{{.Name}}
}

// New{{.Name}}Builder 创建新的构建器
func New{{.Name}}Builder() *{{.Name}}Builder {
    return &{{.Name}}Builder{
        instance: &{{.Name}}{},
    }
}

{{range .Fields}}
// {{.Name}} 设置 {{.Name}} 字段
func (b *{{$.Name}}Builder) {{.Name}}(value {{.Type}}) *{{$.Name}}Builder {
    b.instance.{{.Name}} = value
    return b
}
{{end}}

// Build 构建 {{.Name}} 实例
func (b *{{.Name}}Builder) Build() *{{.Name}} {
    return b.instance
}

{{end}}
`

func main() {
    // 解析当前包的源代码
    fset := token.NewFileSet()
    pkgs, err := parser.ParseDir(fset, ".", nil, parser.ParseComments)
    if err != nil {
        fmt.Printf("Error parsing package: %v\n", err)
        os.Exit(1)
    }

    var structs []StructInfo

    // 遍历包中的所有文件
    for _, pkg := range pkgs {
        for _, file := range pkg.Files {
            // 查找结构体定义
            ast.Inspect(file, func(n ast.Node) bool {
                switch x := n.(type) {
                case *ast.GenDecl:
                    if x.Tok == token.TYPE {
                        for _, spec := range x.Specs {
                            if ts, ok := spec.(*ast.TypeSpec); ok {
                                if st, ok := ts.Type.(*ast.StructType); ok {
                                    structInfo := extractStructInfo(ts.Name.Name, st)
                                    structs = append(structs, structInfo)
                                }
                            }
                        }
                    }
                }
                return true
            })
        }
    }

    // 生成代码
    if len(structs) > 0 {
        generateBuilders(structs)
    }
}

func extractStructInfo(name string, st *ast.StructType) StructInfo {
    var fields []FieldInfo

    for _, field := range st.Fields.List {
        for _, fieldName := range field.Names {
            fieldType := getTypeString(field.Type)
            fieldTag := ""
            if field.Tag != nil {
                fieldTag = field.Tag.Value
            }

            fields = append(fields, FieldInfo{
                Name: fieldName.Name,
                Type: fieldType,
                Tag:  fieldTag,
            })
        }
    }

    return StructInfo{
        Name:   name,
        Fields: fields,
    }
}

func getTypeString(expr ast.Expr) string {
    switch x := expr.(type) {
    case *ast.Ident:
        return x.Name
    case *ast.SelectorExpr:
        return getTypeString(x.X) + "." + x.Sel.Name
    case *ast.StarExpr:
        return "*" + getTypeString(x.X)
    case *ast.ArrayType:
        return "[]" + getTypeString(x.Elt)
    default:
        return "interface{}"
    }
}

func generateBuilders(structs []StructInfo) {
    tmpl, err := template.New("builder").Parse(builderTemplate)
    if err != nil {
        fmt.Printf("Error parsing template: %v\n", err)
        return
    }

    file, err := os.Create("generated_builders.go")
    if err != nil {
        fmt.Printf("Error creating file: %v\n", err)
        return
    }
    defer file.Close()

    err = tmpl.Execute(file, structs)
    if err != nil {
        fmt.Printf("Error executing template: %v\n", err)
        return
    }

    fmt.Println("Generated builders for", len(structs), "structs")
}

package main

import (
    "fmt"
    "go/format"
    "os"
    "reflect"
    "strings"
    "text/template"
)

// FieldDefinition 字段定义
type FieldDefinition struct {
    Name string
    Type string
    Tag  string
}

// StructDefinition 结构体定义
type StructDefinition struct {
    Name   string
    Fields []FieldDefinition
}

// CodeGenerator 代码生成器
type CodeGenerator struct {
    packageName string
    imports     []string
    structs     []StructDefinition
}

// NewCodeGenerator 创建代码生成器
func NewCodeGenerator(packageName string) *CodeGenerator {
    return &CodeGenerator{
        packageName: packageName,
        imports:     []string{},
        structs:     []StructDefinition{},
    }
}

// AddImport 添加导入
func (cg *CodeGenerator) AddImport(importPath string) {
    for _, existing := range cg.imports {
        if existing == importPath {
            return
        }
    }
    cg.imports = append(cg.imports, importPath)
}

// AddStruct 添加结构体定义
func (cg *CodeGenerator) AddStruct(structDef StructDefinition) {
    cg.structs = append(cg.structs, structDef)
}

// GenerateFromReflection 从反射信息生成结构体
func (cg *CodeGenerator) GenerateFromReflection(obj interface{}) {
    t := reflect.TypeOf(obj)
    if t.Kind() == reflect.Ptr {
        t = t.Elem()
    }

    if t.Kind() != reflect.Struct {
        return
    }

    structDef := StructDefinition{
        Name:   t.Name(),
        Fields: []FieldDefinition{},
    }

    for i := 0; i < t.NumField(); i++ {
        field := t.Field(i)

        fieldDef := FieldDefinition{
            Name: field.Name,
            Type: cg.getTypeString(field.Type),
            Tag:  string(field.Tag),
        }

        structDef.Fields = append(structDef.Fields, fieldDef)
    }

    cg.AddStruct(structDef)
}

// getTypeString 获取类型字符串
func (cg *CodeGenerator) getTypeString(t reflect.Type) string {
    switch t.Kind() {
    case reflect.Ptr:
        return "*" + cg.getTypeString(t.Elem())
    case reflect.Slice:
        return "[]" + cg.getTypeString(t.Elem())
    case reflect.Array:
        return fmt.Sprintf("[%d]%s", t.Len(), cg.getTypeString(t.Elem()))
    case reflect.Map:
        return fmt.Sprintf("map[%s]%s",
            cg.getTypeString(t.Key()),
            cg.getTypeString(t.Elem()))
    case reflect.Chan:
        return "chan " + cg.getTypeString(t.Elem())
    case reflect.Func:
        return cg.getFuncTypeString(t)
    default:
        if t.PkgPath() != "" && t.PkgPath() != cg.packageName {
            // 添加导入
            cg.AddImport(t.PkgPath())
            parts := strings.Split(t.PkgPath(), "/")
            pkgName := parts[len(parts)-1]
            return pkgName + "." + t.Name()
        }
        return t.Name()
    }
}

// getFuncTypeString 获取函数类型字符串
func (cg *CodeGenerator) getFuncTypeString(t reflect.Type) string {
    var parts []string
    parts = append(parts, "func(")

    // 输入参数
    var inParams []string
    for i := 0; i < t.NumIn(); i++ {
        inParams = append(inParams, cg.getTypeString(t.In(i)))
    }
    parts = append(parts, strings.Join(inParams, ", "))
    parts = append(parts, ")")

    // 输出参数
    if t.NumOut() > 0 {
        if t.NumOut() == 1 {
            parts = append(parts, " "+cg.getTypeString(t.Out(0)))
        } else {
            var outParams []string
            for i := 0; i < t.NumOut(); i++ {
                outParams = append(outParams, cg.getTypeString(t.Out(i)))
            }
            parts = append(parts, " ("+strings.Join(outParams, ", ")+")")
        }
    }

    return strings.Join(parts, "")
}

// Generate 生成代码
func (cg *CodeGenerator) Generate() (string, error) {
    tmpl := `package {{.PackageName}}

{{if .Imports}}
import (
{{range .Imports}}    "{{.}}"
{{end}})
{{end}}

{{range .Structs}}
type {{.Name}} struct {
{{range .Fields}}    {{.Name}} {{.Type}}{{if .Tag}} ` + "`{{.Tag}}`" + `{{end}}
{{end}}}

{{end}}`

    t, err := template.New("code").Parse(tmpl)
    if err != nil {
        return "", err
    }

    var buf strings.Builder
    err = t.Execute(&buf, map[string]interface{}{
        "PackageName": cg.packageName,
        "Imports":     cg.imports,
        "Structs":     cg.structs,
    })
    if err != nil {
        return "", err
    }

    // 格式化代码
    formatted, err := format.Source([]byte(buf.String()))
    if err != nil {
        return buf.String(), nil // 返回未格式化的代码
    }

    return string(formatted), nil
}

// SaveToFile 保存到文件
func (cg *CodeGenerator) SaveToFile(filename string) error {
    code, err := cg.Generate()
    if err != nil {
        return err
    }

    return os.WriteFile(filename, []byte(code), 0644)
}

// 示例结构体
type User struct {
    ID       int       `json:"id" db:"user_id"`
    Name     string    `json:"name" db:"username"`
    Email    string    `json:"email" db:"email"`
    Settings *Settings `json:"settings,omitempty"`
}

type Settings struct {
    Theme    string            `json:"theme"`
    Language string            `json:"language"`
    Metadata map[string]string `json:"metadata"`
}

func main() {
    generator := NewCodeGenerator("generated")

    // 从现有结构体生成代码
    generator.GenerateFromReflection(User{})
    generator.GenerateFromReflection(Settings{})

    // 手动添加结构体
    generator.AddStruct(StructDefinition{
        Name: "Product",
        Fields: []FieldDefinition{
            {Name: "ID", Type: "int", Tag: `json:"id"`},
            {Name: "Name", Type: "string", Tag: `json:"name"`},
            {Name: "Price", Type: "float64", Tag: `json:"price"`},
            {Name: "InStock", Type: "bool", Tag: `json:"in_stock"`},
        },
    })

    // 生成代码
    code, err := generator.Generate()
    if err != nil {
        fmt.Printf("Error generating code: %v\n", err)
        return
    }

    fmt.Println("Generated code:")
    fmt.Println(code)

    // 保存到文件
    err = generator.SaveToFile("generated_structs.go")
    if err != nil {
        fmt.Printf("Error saving file: %v\n", err)
    } else {
        fmt.Println("Code saved to generated_structs.go")
    }
}

package main

import (
    "fmt"
    "os"
    "strings"
    "text/template"
    "time"
)

// TemplateData 模板数据
type TemplateData struct {
    PackageName   string
    GeneratedTime string
    Author        string
    Structs       []StructTemplate
    Functions     []FunctionTemplate
}

// StructTemplate 结构体模板
type StructTemplate struct {
    Name        string
    Comment     string
    Fields      []FieldTemplate
    Methods     []MethodTemplate
    Implements  []string
    Tags        map[string]string
}

// FieldTemplate 字段模板
type FieldTemplate struct {
    Name    string
    Type    string
    Tag     string
    Comment string
    Default string
}

// MethodTemplate 方法模板
type MethodTemplate struct {
    Name       string
    Receiver   string
    Parameters []ParameterTemplate
    Returns    []ParameterTemplate
    Body       string
    Comment    string
}

// FunctionTemplate 函数模板
type FunctionTemplate struct {
    Name       string
    Parameters []ParameterTemplate
    Returns    []ParameterTemplate
    Body       string
    Comment    string
}

// ParameterTemplate 参数模板
type ParameterTemplate struct {
    Name string
    Type string
}

// AdvancedGenerator 高级代码生成器
type AdvancedGenerator struct {
    templates map[string]*template.Template
    funcMap   template.FuncMap
}

// NewAdvancedGenerator 创建高级代码生成器
func NewAdvancedGenerator() *AdvancedGenerator {
    ag := &AdvancedGenerator{
        templates: make(map[string]*template.Template),
        funcMap: template.FuncMap{
            "title":      strings.Title,
            "lower":      strings.ToLower,
            "upper":      strings.ToUpper,
            "camelCase":  toCamelCase,
            "snakeCase":  toSnakeCase,
            "pluralize":  pluralize,
            "join":       strings.Join,
            "hasPrefix":  strings.HasPrefix,
            "hasSuffix":  strings.HasSuffix,
            "contains":   strings.Contains,
            "replace":    strings.ReplaceAll,
            "now":        time.Now,
            "formatTime": formatTime,
        },
    }

    ag.loadDefaultTemplates()
    return ag
}

// 辅助函数
func toCamelCase(s string) string {
    parts := strings.Split(s, "_")
    for i := 1; i < len(parts); i++ {
        parts[i] = strings.Title(parts[i])
    }
    return strings.Join(parts, "")
}

func toSnakeCase(s string) string {
    var result []rune
    for i, r := range s {
        if i > 0 && r >= 'A' && r <= 'Z' {
            result = append(result, '_')
        }
        result = append(result, r)
    }
    return strings.ToLower(string(result))
}

func pluralize(s string) string {
    if strings.HasSuffix(s, "y") {
        return s[:len(s)-1] + "ies"
    }
    if strings.HasSuffix(s, "s") || strings.HasSuffix(s, "x") || strings.HasSuffix(s, "z") {
        return s + "es"
    }
    return s + "s"
}

func formatTime(t time.Time, layout string) string {
    return t.Format(layout)
}

// loadDefaultTemplates 加载默认模板
func (ag *AdvancedGenerator) loadDefaultTemplates() {
    // CRUD 操作模板
    crudTemplate := `// Code generated by AdvancedGenerator at {{formatTime .GeneratedTime "2006-01-02 15:04:05"}}
// Author: {{.Author}}
// DO NOT EDIT.

package {{.PackageName}}

import (
    "database/sql"
    "fmt"
)

{{range .Structs}}
{{if .Comment}}// {{.Comment}}{{end}}
type {{.Name}} struct {
{{range .Fields}}    {{.Name}} {{.Type}}{{if .Tag}} ` + "`{{.Tag}}`" + `{{end}}{{if .Comment}} // {{.Comment}}{{end}}
{{end}}}

// {{.Name}}Repository 提供 {{.Name}} 的数据库操作
type {{.Name}}Repository struct {
    db *sql.DB
}

// New{{.Name}}Repository 创建新的仓库实例
func New{{.Name}}Repository(db *sql.DB) *{{.Name}}Repository {
    return &{{.Name}}Repository{db: db}
}

// Create 创建新的 {{.Name}}
func (r *{{.Name}}Repository) Create({{lower .Name}} *{{.Name}}) error {
    query := ` + "`INSERT INTO {{snakeCase .Name | pluralize}} ({{range $i, $f := .Fields}}{{if $i}}, {{end}}{{snakeCase $f.Name}}{{end}}) VALUES ({{range $i, $f := .Fields}}{{if $i}}, {{end}}?{{end}})`" + `
    _, err := r.db.Exec(query{{range .Fields}}, {{lower $.Name}}.{{.Name}}{{end}})
    return err
}

// GetByID 根据ID获取 {{.Name}}
func (r *{{.Name}}Repository) GetByID(id int) (*{{.Name}}, error) {
    {{lower .Name}} := &{{.Name}}{}
    query := ` + "`SELECT {{range $i, $f := .Fields}}{{if $i}}, {{end}}{{snakeCase $f.Name}}{{end}} FROM {{snakeCase .Name | pluralize}} WHERE id = ?`" + `
    err := r.db.QueryRow(query, id).Scan({{range $i, $f := .Fields}}{{if $i}}, {{end}}&{{lower $.Name}}.{{$f.Name}}{{end}})
    if err != nil {
        return nil, err
    }
    return {{lower .Name}}, nil
}

// Update 更新 {{.Name}}
func (r *{{.Name}}Repository) Update({{lower .Name}} *{{.Name}}) error {
    query := ` + "`UPDATE {{snakeCase .Name | pluralize}} SET {{range $i, $f := .Fields}}{{if $i}}, {{end}}{{snakeCase $f.Name}} = ?{{end}} WHERE id = ?`" + `
    _, err := r.db.Exec(query{{range .Fields}}, {{lower $.Name}}.{{.Name}}{{end}}, {{lower .Name}}.ID)
    return err
}

// Delete 删除 {{.Name}}
func (r *{{.Name}}Repository) Delete(id int) error {
    query := ` + "`DELETE FROM {{snakeCase .Name | pluralize}} WHERE id = ?`" + `
    _, err := r.db.Exec(query, id)
    return err
}

// List 获取所有 {{.Name}}
func (r *{{.Name}}Repository) List() ([]*{{.Name}}, error) {
    query := ` + "`SELECT {{range $i, $f := .Fields}}{{if $i}}, {{end}}{{snakeCase $f.Name}}{{end}} FROM {{snakeCase .Name | pluralize}}`" + `
    rows, err := r.db.Query(query)
    if err != nil {
        return nil, err
    }
    defer rows.Close()

    var {{pluralize (lower .Name)}} []*{{.Name}}
    for rows.Next() {
        {{lower .Name}} := &{{.Name}}{}
        err := rows.Scan({{range $i, $f := .Fields}}{{if $i}}, {{end}}&{{lower $.Name}}.{{$f.Name}}{{end}})
        if err != nil {
            return nil, err
        }
        {{pluralize (lower .Name)}} = append({{pluralize (lower .Name)}}, {{lower .Name}})
    }

    return {{pluralize (lower .Name)}}, nil
}

{{end}}`

    ag.AddTemplate("crud", crudTemplate)
}

// AddTemplate 添加模板
func (ag *AdvancedGenerator) AddTemplate(name, templateStr string) error {
    tmpl, err := template.New(name).Funcs(ag.funcMap).Parse(templateStr)
    if err != nil {
        return err
    }
    ag.templates[name] = tmpl
    return nil
}

// Generate 生成代码
func (ag *AdvancedGenerator) Generate(templateName string, data TemplateData) (string, error) {
    tmpl, exists := ag.templates[templateName]
    if !exists {
        return "", fmt.Errorf("template '%s' not found", templateName)
    }

    var buf strings.Builder
    err := tmpl.Execute(&buf, data)
    if err != nil {
        return "", err
    }

    return buf.String(), nil
}

// GenerateToFile 生成代码到文件
func (ag *AdvancedGenerator) GenerateToFile(templateName string, data TemplateData, filename string) error {
    code, err := ag.Generate(templateName, data)
    if err != nil {
        return err
    }

    return os.WriteFile(filename, []byte(code), 0644)
}

func main() {
    generator := NewAdvancedGenerator()

    // 准备模板数据
    data := TemplateData{
        PackageName:   "models",
        GeneratedTime: time.Now().Format(time.RFC3339),
        Author:        "Code Generator",
        Structs: []StructTemplate{
            {
                Name:    "User",
                Comment: "User 表示系统用户",
                Fields: []FieldTemplate{
                    {Name: "ID", Type: "int", Tag: `json:"id" db:"id"`, Comment: "用户ID"},
                    {Name: "Username", Type: "string", Tag: `json:"username" db:"username"`, Comment: "用户名"},
                    {Name: "Email", Type: "string", Tag: `json:"email" db:"email"`, Comment: "邮箱地址"},
                    {Name: "CreatedAt", Type: "time.Time", Tag: `json:"created_at" db:"created_at"`, Comment: "创建时间"},
                },
            },
            {
                Name:    "Product",
                Comment: "Product 表示商品信息",
                Fields: []FieldTemplate{
                    {Name: "ID", Type: "int", Tag: `json:"id" db:"id"`, Comment: "商品ID"},
                    {Name: "Name", Type: "string", Tag: `json:"name" db:"name"`, Comment: "商品名称"},
                    {Name: "Price", Type: "float64", Tag: `json:"price" db:"price"`, Comment: "价格"},
                    {Name: "Stock", Type: "int", Tag: `json:"stock" db:"stock"`, Comment: "库存数量"},
                },
            },
        },
    }

    // 生成 CRUD 代码
    fmt.Println("Generating CRUD code...")
    err := generator.GenerateToFile("crud", data, "generated_crud.go")
    if err != nil {
        fmt.Printf("Error generating CRUD code: %v\n", err)
    } else {
        fmt.Println("CRUD code generated successfully!")
    }

    // 显示生成的代码片段
    crudCode, _ := generator.Generate("crud", data)
    fmt.Println("\nGenerated CRUD code preview:")
    fmt.Println(crudCode[:500] + "...")
}