Golang AST语法树使用教程及示例
许多自动化代码生成工具都离不开语法树分析,例如 goimport
, gomock
, wire
等项目都离不开语法树分析。基于语法树分析,可以实现许多有趣实用的工具。本篇将结合示例,展示如何基于 ast
标准包操作语法树。
本篇中的代码的完整示例可以在这里找到: ast-example
Quick Start
首先我们看下语法树长什么样子,以下代码将打印 ./demo.go
文件的语法树:
package main import ( "go/ast" "go/parser" "go/token" "log" "path/filepath" ) func main() { fset := token.NewFileSet() // 这里取绝对路径,方便打印出来的语法树可以转跳到编辑器 path, _ := filepath.Abs("./demo.go") f, err := parser.ParseFile(fset, path, nil, parser.AllErrors) if err != nil { log.Println(err) return } // 打印语法树 ast.Print(fset, f) } 复制代码
demo.go:
package main import ( "context" ) // Foo 结构体 type Foo struct { i int } // Bar 接口 type Bar interface { Do(ctx context.Context) error } // main方法 func main() { a := 1 } 复制代码
demo.go
文件已尽量简化,但其语法树的输出内容依旧十分庞大。我们截取部分来做一些简要的说明。
首先是文件所属的包名,和其声明在文件中的位置:
0 *ast.File { 1 . Package: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:1 2 . Name: *ast.Ident { 3 . . NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:9 4 . . Name: "main" 5 . } ... 复制代码
紧接着是 Decls
,也就是Declarations,其包含了声明的一些变量,方法,接口等:
... 6 . Decls: []ast.Decl (len = 4) { 7 . . 0: *ast.GenDecl { 8 . . . TokPos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:3:1 9 . . . Tok: import 10 . . . Lparen: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:3:8 11 . . . Specs: []ast.Spec (len = 1) { 12 . . . . 0: *ast.ImportSpec { 13 . . . . . Path: *ast.BasicLit { 14 . . . . . . ValuePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:4:2 15 . . . . . . Kind: STRING 16 . . . . . . Value: "\"context\"" 17 . . . . . } 18 . . . . . EndPos: - 19 . . . . } 20 . . . } 21 . . . Rparen: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:5:1 22 . . } .... 复制代码
可以看到该语法树包含了4条 Decl
记录,我们取第一条记录为例,该记录为 *ast.GenDecl
类型。不难看出这条记录对应的是我们的 import
代码段。始位置(TokPos),左右括号的位置(Lparen,Rparen),和import的包(Specs)等信息都能从语法树中得到。
语法树的打印信来自 ast.File
结构体:
$GOROOT/src/go/ast/ast.go
// 该结构体位于标准包 go/ast/ast.go 中,有兴趣可以转跳到源码阅读更详尽的注释 type File struct { Doc *CommentGroup // associated documentation; or nil Package token.Pos // position of "package" keyword Name *Ident // package name Decls []Decl // top-level declarations; or nil Scope *Scope // package scope (this file only) Imports []*ImportSpec // imports in this file Unresolved []*Ident // unresolved identifiers in this file Comments []*CommentGroup // list of all comments in the source file } 复制代码
结合注释和字段名我们大概知道每个字段的含义,接下来我们详细梳理一下语法树的组成结构。
Node节点
整个语法树由不同的node组成,从源码注释中可以得知主要有如下三种node:
There are 3 main classes of nodes: Expressions and type nodes, statement nodes, and declaration nodes.
在Go的 Language Specification
中可以找到这些节点类型详细规范和说明,有兴趣的小伙伴可以深入研究一下,在此不做展开。
但实际在代码,出现了第四种node:Spec Node,每种node都有专门的接口定义:
$GOROOT/src/go/ast/ast.go
... // All node types implement the Node interface. type Node interface { Pos() token.Pos // position of first character belonging to the node End() token.Pos // position of first character immediately after the node } // All expression nodes implement the Expr interface. type Expr interface { Node exprNode() } // All statement nodes implement the Stmt interface. type Stmt interface { Node stmtNode() } // All declaration nodes implement the Decl interface. type Decl interface { Node declNode() } ... // A Spec node represents a single (non-parenthesized) import, // constant, type, or variable declaration. // type ( // The Spec type stands for any of *ImportSpec, *ValueSpec, and *TypeSpec. Spec interface { Node specNode() } .... ) 复制代码
可以看到所有的node都继承 Node
接口,记录了node的开始和结束位置。还记得Quick Start示例中的 Decls
吗?它正是 declaration nodes
。除去上述四种使用接口进行分类的node,还有些node没有再额外定义接口细分类别,仅实现了 Node
接口,为了方便描述,在本篇中我把这些节点称为 common node
。 $GOROOT/src/go/ast/ast.go
列举了所有所有节点的实现,我们从中挑选几个作为例子,感受一下它们的区别。
Expression and Type
先来看expression node。
$GOROOT/src/go/ast/ast.go
... // An Ident node represents an identifier. Ident struct { NamePos token.Pos // identifier position Name string // identifier name Obj *Object // denoted object; or nil } ... 复制代码
Indent
(identifier)表示一个标识符,比如Quick Start示例中表示包名的 Name
字段就是一个expression node:
0 *ast.File { 1 . Package: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:1 2 . Name: *ast.Ident { <---- 3 . . NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:9 4 . . Name: "main" 5 . } ... 复制代码
接下来是type node。
$GOROOT/src/go/ast/ast.go
... // A StructType node represents a struct type. StructType struct { Struct token.Pos // position of "struct" keyword Fields *FieldList // list of field declarations Incomplete bool // true if (source) fields are missing in the Fields list } // Pointer types are represented via StarExpr nodes. // A FuncType node represents a function type. FuncType struct { Func token.Pos // position of "func" keyword (token.NoPos if there is no "func") Params *FieldList // (incoming) parameters; non-nil Results *FieldList // (outgoing) results; or nil } // An InterfaceType node represents an interface type. InterfaceType struct { Interface token.Pos // position of "interface" keyword Methods *FieldList // list of methods Incomplete bool // true if (source) methods are missing in the Methods list } ... 复制代码
type node很好理解,它包含一些复合类型,例如在Quick Start中出现的 StructType
, FuncType
和 InterfaceType
。
Statement
赋值语句,控制语句(if,else,for,select…)等均属于statement node。
$GOROOT/src/go/ast/ast.go
... // An AssignStmt node represents an assignment or // a short variable declaration. // AssignStmt struct { Lhs []Expr TokPos token.Pos // position of Tok Tok token.Token // assignment token, DEFINE Rhs []Expr } ... // An IfStmt node represents an if statement. IfStmt struct { If token.Pos // position of "if" keyword Init Stmt // initialization statement; or nil Cond Expr // condition Body *BlockStmt Else Stmt // else branch; or nil } ... 复制代码
例如Quick Start中,我们在 main
函数中对变量a赋值的程序片段就属于 AssignStmt
:
... 174 . . . Body: *ast.BlockStmt { 175 . . . . Lbrace: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:18:13 176 . . . . List: []ast.Stmt (len = 1) { 177 . . . . . 0: *ast.AssignStmt { <--- 这里 178 . . . . . . Lhs: []ast.Expr (len = 1) { 179 . . . . . . . 0: *ast.Ident { 180 . . . . . . . . NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:19:2 181 . . . . . . . . Name: "a" ... 复制代码
Spec Node
Spec node只有3种,分别是 ImportSpec
, ValueSpec
和 TypeSpec
:
$GOROOT/src/go/ast/ast.go
// An ImportSpec node represents a single package import. ImportSpec struct { Doc *CommentGroup // associated documentation; or nil Name *Ident // local package name (including "."); or nil Path *BasicLit // import path Comment *CommentGroup // line comments; or nil EndPos token.Pos // end of spec (overrides Path.Pos if nonzero) } // A ValueSpec node represents a constant or variable declaration // (ConstSpec or VarSpec production). // ValueSpec struct { Doc *CommentGroup // associated documentation; or nil Names []*Ident // value names (len(Names) > 0) Type Expr // value type; or nil Values []Expr // initial values; or nil Comment *CommentGroup // line comments; or nil } // A TypeSpec node represents a type declaration (TypeSpec production). TypeSpec struct { Doc *CommentGroup // associated documentation; or nil Name *Ident // type name Assign token.Pos // position of '=', if any Type Expr // *Ident, *ParenExpr, *SelectorExpr, *StarExpr, or any of the *XxxTypes Comment *CommentGroup // line comments; or nil } 复制代码
ImportSpec
表示一个单独的import, ValueSpec
表示一个常量或变量的声明, TypeSpec
则表示一个type声明。例如
在Quick Start示例中,出现了 ImportSpec
和 TypeSpec
import ( "context" // <--- 这里是一个ImportSpec node ) // Foo 结构体 type Foo struct { // <--- 这里是一个TypeSpec node i int } 复制代码
在语法树的打印结果中可以看到对应的输出,小伙伴们可自行查找。
Declaration Node
Declaration node也只有三种:
$GOROOT/src/go/ast/ast.go
... type ( // A BadDecl node is a placeholder for declarations containing // syntax errors for which no correct declaration nodes can be // created. // BadDecl struct { From, To token.Pos // position range of bad declaration } // A GenDecl node (generic declaration node) represents an import, // constant, type or variable declaration. A valid Lparen position // (Lparen.IsValid()) indicates a parenthesized declaration. // // Relationship between Tok value and Specs element type: // // token.IMPORT *ImportSpec // token.CONST *ValueSpec // token.TYPE *TypeSpec // token.VAR *ValueSpec // GenDecl struct { Doc *CommentGroup // associated documentation; or nil TokPos token.Pos // position of Tok Tok token.Token // IMPORT, CONST, TYPE, VAR Lparen token.Pos // position of '(', if any Specs []Spec Rparen token.Pos // position of ')', if any } // A FuncDecl node represents a function declaration. FuncDecl struct { Doc *CommentGroup // associated documentation; or nil Recv *FieldList // receiver (methods); or nil (functions) Name *Ident // function/method name Type *FuncType // function signature: parameters, results, and position of "func" keyword Body *BlockStmt // function body; or nil for external (non-Go) function } ) ... 复制代码
BadDecl
表示一个有语法错误的节点; GenDecl
用于表示import, const,type或变量声明; FunDecl
用于表示函数声明。 GenDecl
和 FunDecl
在Quick Start例子中均有出现,小伙伴们可自行查找。
Common Node
除去上述四种类别划分的node,还有一些node不属于上面四种类别:
$GOROOT/src/go/ast/ast.go
// Comment 注释节点,代表单行的 //-格式 或 /*-格式的注释. type Comment struct { ... } ... // CommentGroup 注释块节点,包含多个连续的Comment type CommentGroup struct { ... } // Field 字段节点, 可以代表结构体定义中的字段,接口定义中的方法列表,函数前面中的入参和返回值字段 type Field struct { ... } ... // FieldList 包含多个Field type FieldList struct { ... } // File 表示一个文件节点 type File struct { ... } // Package 表示一个包节点 type Package struct { ... } 复制代码
Quick Start示例包含了上面列举的所有node,小伙伴们可以自行查找。更为详细的注释和具体的结构体字段请查阅源码。
所有的节点类型大致列举完毕,其中还有许多具体的节点类型未能一一列举,但基本上都是大同小异,源码注释也比较清晰,等用到的时候再细看也不迟。现在我们对整个语法树的构造有了基本的了解,接下来通过几个示例来演示具体用法。
示例
为文件中所有接口方法添加context参数
实现这个功能我们需要四步:
context context.Context
遍历语法树
语法树层级较深,嵌套关系复杂,如果不能完全掌握node之间的关系和嵌套规则,我们很难自己写出正确的遍历方法。不过好在 ast
包已经为我们提供了遍历方法:
$GOROOT/src/go/ast/ast.go
func Walk(v Visitor, node Node) 复制代码
type Visitor interface { Visit(node Node) (w Visitor) } 复制代码
Walk
方法会按照深度优先搜索方法(depth-first order)遍历整个语法树,我们只需按照我们的业务需要,实现 Visitor
接口即可。 Walk
每遍历一个节点就会调用 Visitor.Visit
方法,传入当前节点。如果 Visit
返回 nil
,则停止遍历当前节点的子节点。本示例的 Visitor
实现如下:
// Visitor type Visitor struct { } func (v *Visitor) Visit(node ast.Node) ast.Visitor { switch node.(type) { case *ast.GenDecl: genDecl := node.(*ast.GenDecl) // 查找有没有import context包 // Notice:没有考虑没有import任何包的情况 if genDecl.Tok == token.IMPORT { v.addImport(genDecl) // 不需要再遍历子树 return nil } case *ast.InterfaceType: // 遍历所有的接口类型 iface := node.(*ast.InterfaceType) addContext(iface) // 不需要再遍历子树 return nil } return v } 复制代码
添加import
// addImport 引入context包 func (v *Visitor) addImport(genDecl *ast.GenDecl) { // 是否已经import hasImported := false for _, v := range genDecl.Specs { imptSpec := v.(*ast.ImportSpec) // 如果已经包含"context" if imptSpec.Path.Value == strconv.Quote("context") { hasImported = true } } // 如果没有import context,则import if !hasImported { genDecl.Specs = append(genDecl.Specs, &ast.ImportSpec{ Path: &ast.BasicLit{ Kind: token.STRING, Value: strconv.Quote("context"), }, }) } } 复制代码
为接口方法添加参数
// addContext 添加context参数 func addContext(iface *ast.InterfaceType) { // 接口方法不为空时,遍历接口方法 if iface.Methods != nil || iface.Methods.List != nil { for _, v := range iface.Methods.List { ft := v.Type.(*ast.FuncType) hasContext := false // 判断参数中是否包含context.Context类型 for _, v := range ft.Params.List { if expr, ok := v.Type.(*ast.SelectorExpr); ok { if ident, ok := expr.X.(*ast.Ident); ok { if ident.Name == "context" { hasContext = true } } } } // 为没有context参数的方法添加context参数 if !hasContext { ctxField := &ast.Field{ Names: []*ast.Ident{ ast.NewIdent("ctx"), }, // Notice: 没有考虑import别名的情况 Type: &ast.SelectorExpr{ X: ast.NewIdent("context"), Sel: ast.NewIdent("Context"), }, } list := []*ast.Field{ ctxField, } ft.Params.List = append(list, ft.Params.List...) } } } } 复制代码
将语法树转换成Go代码
format
包为我们提供了转换函数, format.Node
会将语法树按照 gofmt
的格式输出:
... var output []byte buffer := bytes.NewBuffer(output) err = format.Node(buffer, fset, f) if err != nil { log.Fatal(err) } // 输出Go代码 fmt.Println(buffer.String()) ... 复制代码
输出结果如下:
package main import ( "context" ) type Foo interface { FooA(ctx context.Context, i int) FooB(ctx context.Context, j int) FooC(ctx context.Context) } type Bar interface { BarA(ctx context.Context, i int) BarB(ctx context.Context) BarC(ctx context.Context) } 复制代码
可以看到我们所有的接口方的第一个参数都变成了 context.Context
。建议将示例中的语法树先打印出来,再对照着代码看,方便理解。
一些坑与不足
至此我们已经完成了语法树的解析,遍历,修改以及输出。但细心的小伙伴可能已经发现:示例中的文件并没有出现一行注释。这的确是有意为之,如果我们加上注释,会发现最终生成文件的注释就像迷途的羔羊,完全找不到自己的位置。比如这样:
//修改前 type Foo interface { FooA(i int) // FooB FooB(j int) FooC(ctx context.Context) } // 修改后 type Foo interface { FooA(ctx context. // FooB Context, i int) FooB(ctx context.Context, j int) FooC(ctx context.Context) } 复制代码
导致这种现象的原因在于: ast
包生成的语法树中的注释是”free-floating”的。还记得每个node都有 Pos()
和 End()
方法来标识其位置吗?对于非注释节点,语法树能够正确的调整他们的位置,但却不能自动调整注释节点的位置。如果我们想要让注释出现在正确的位置上,我们必须手动设置节点 Pos
和 End
。源码注释中提到了这个问题:
Whether and how a comment is associated with a node depends on the interpretation of the syntax tree by the manipulating program: Except for Doc and Comment comments directly associated with nodes, the remaining comments are "free-floating" (see also issues #18593, #20744).
issue中有具体的讨论,官方承认这是一个设计缺陷,但还是迟迟未能改进。其中有位迫不及待的小哥提供了自己的方案:
如果实在是要对有注释的语法树进行修改,可以尝试一下。
虽然语法树的确存在修改困难问题,但其还是能满足大部分基于语法树分析的代码生成工作了(gomock,wire等等)。