文章

Go AST 与工具链

Go AST 与工具链

包总览

核心类型用途
go/astNode, Decl, Expr, StmtAST 节点定义
go/parserParseFile, ParseDir源码解析
go/tokenFileSet, Position位置信息
go/formatNode, Source代码格式化
go/typesChecker, Info类型检查
go/importerImport包导入
go/printerFprintAST 打印
go/buildPackage构建信息

这些包用于代码分析、代码生成、静态检查等场景。日常开发很少直接使用,但对编写工具链非常重要。

go/parser + go/ast 基本用法

解析源码

import (
    "go/ast"
    "go/parser"
    "go/token"
)

// 解析源码文件
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, "main.go", nil, parser.ParseComments)
// 参数:
//   fset      - 位置信息集合
//   filename  - 文件名(空字符串则用 src)
//   src       - 源码内容(nil 则从文件读取)
//   mode      - 解析模式

// 解析源码字符串
f, err := parser.ParseFile(fset, "", src, parser.ParseComments)

// 解析目录下所有文件
pkgs, err := parser.ParseDir(fset, "./mypackage", nil, parser.ParseComments)

遍历 AST

src := `
package main

import "fmt"

type User struct {
    Name string
    Age  int
}

func main() {
    u := User{Name: "小徐", Age: 28}
    fmt.Println(u.Name)
}
`

fset := token.NewFileSet()
f, _ := parser.ParseFile(fset, "", src, parser.ParseComments)

// 遍历所有节点
ast.Inspect(f, func(n ast.Node) bool {
    switch x := n.(type) {
    case *ast.File:
        fmt.Println("Package:", x.Name.Name)
    case *ast.ImportSpec:
        fmt.Println("Import:", x.Path.Value)
    case *ast.TypeSpec:
        fmt.Println("Type:", x.Name.Name)
    case *ast.FuncDecl:
        fmt.Println("Func:", x.Name.Name)
    case *ast.StructType:
        for _, field := range x.Fields.List {
            for _, name := range field.Names {
                fmt.Printf("  Field: %s\n", name.Name)
            }
        }
    }
    return true
})

提取所有函数

func ListFunctions(f *ast.File) []string {
    var funcs []string
    for _, decl := range f.Decls {
        if fn, ok := decl.(*ast.FuncDecl); ok {
            receiver := ""
            if fn.Recv != nil {
                // 方法接收者
                if expr, ok := fn.Recv.List[0].Type.(*ast.Ident); ok {
                    receiver = expr.Name + "."
                }
            }
            funcs = append(funcs, receiver+fn.Name.Name)
        }
    }
    return funcs
}

提取 struct 及其字段

func ListStructs(f *ast.File) map[string][]string {
    structs := make(map[string][]string)
    for _, decl := range f.Decls {
        genDecl, ok := decl.(*ast.GenDecl)
        if !ok || genDecl.Tok != token.TYPE {
            continue
        }
        for _, spec := range genDecl.Specs {
            typeSpec := spec.(*ast.TypeSpec)
            structType, ok := typeSpec.Type.(*ast.StructType)
            if !ok {
                continue
            }
            var fields []string
            for _, field := range structType.Fields.List {
                for _, name := range field.Names {
                    fields = append(fields, name.Name)
                }
            }
            structs[typeSpec.Name.Name] = fields
        }
    }
    return structs
}

提取 struct tag

func GetStructTags(f *ast.File) map[string]map[string]string {
    result := make(map[string]map[string]string)

    ast.Inspect(f, func(n ast.Node) bool {
        typeSpec, ok := n.(*ast.TypeSpec)
        if !ok {
            return true
        }
        structType, ok := typeSpec.Type.(*ast.StructType)
        if !ok {
            return true
        }

        tags := make(map[string]string)
        for _, field := range structType.Fields.List {
            if field.Tag == nil {
                continue
            }
            tagValue := field.Tag.Value // 包含反引号
            for _, name := range field.Names {
                tags[name.Name] = tagValue
            }
        }
        result[typeSpec.Name.Name] = tags
        return true
    })
    return result
}

go/format 包

import "go/format"

// 格式化源码
src := `package main
func main(){fmt.Println("hello")}`

formatted, err := format.Source([]byte(src))
// 输出格式化后的代码

// 格式化 AST 节点
fset := token.NewFileSet()
f, _ := parser.ParseFile(fset, "", src, 0)
// 修改 AST...
format.Node(os.Stdout, fset, f)  // 输出格式化后的代码

代码生成示例

生成接口 Mock

import (
    "go/ast"
    "go/parser"
    "go/printer"
    "go/token"
    "os"
)

// 分析接口定义,生成 Mock 代码
func GenerateMock(interfaceFile, interfaceName string) {
    fset := token.NewFileSet()
    f, _ := parser.ParseFile(fset, interfaceFile, nil, 0)

    // 查找接口定义
    var methods []*ast.Field
    for _, decl := range f.Decls {
        genDecl, ok := decl.(*ast.GenDecl)
        if !ok {
            continue
        }
        for _, spec := range genDecl.Specs {
            typeSpec, ok := spec.(*ast.TypeSpec)
            if !ok || typeSpec.Name.Name != interfaceName {
                continue
            }
            iface, ok := typeSpec.Type.(*ast.InterfaceType)
            if !ok {
                continue
            }
            methods = iface.Methods.List
        }
    }

    // 生成 Mock 代码
    fmt.Println("package mock")
    fmt.Printf("type Mock%s struct{}\n", interfaceName)
    for _, m := range methods {
        // 生成每个方法的 mock 实现
        // ...
    }
}

go/types 包(类型检查)

import (
    "go/importer"
    "go/parser"
    "go/token"
    "go/types"
)

// 类型检查
fset := token.NewFileSet()
f, _ := parser.ParseFile(fset, "main.go", nil, 0)

conf := types.Config{
    Importer: importer.Default(),
}

info := &types.Info{
    Types: make(map[ast.Expr]types.TypeAndValue),
    Defs:  make(map[*ast.Ident]types.Object),
    Uses:  make(map[*ast.Ident]types.Object),
}

pkg, err := conf.Check("main", fset, []*ast.File{f}, info)
// 现在可以查询类型信息
// info.Types[expr].Type  - 任意表达式的类型
// info.Defs[ident]       - 标识符定义
// info.Uses[ident]       - 标识符使用

go/build 包

import "go/build"

// 获取包信息
pkg, err := build.Import("github.com/user/project", ".", 0)
pkg.Name          // 包名
pkg.Dir           // 目录路径
pkg.GoFiles       // Go 文件列表
pkg.Imports       // 依赖包列表
pkg.TestGoFiles   // 测试文件列表
pkg.CgoFiles      // CGO 文件列表

实际应用场景

场景工具使用的包
生成 Mock 代码mockgengo/ast, go/parser
生成 String() 方法stringergo/ast, go/types
依赖注入代码生成wirego/ast, go/types
静态检查golangci-lintgo/ast, go/types
代码文档生成goldocgo/ast, go/doc
调用图分析callgraphgo/ast, go/types
ORM 模型生成gengo/ast, go/parser

AST 工具链推荐

用途
golang.org/x/tools/go/packages更强大的包加载(替代 go/build)
golang.org/x/tools/go/ast/astutilAST 修改工具(AddImport 等)
golang.org/x/tools/go/callgraph调用图分析
dave/jennifer代码生成(不基于 AST,更易用)