Golang AST语法树使用教程及示例
source link: https://xiusin.github.io/post/golang-ast-yu-fa-shu-shi-yong-jiao-cheng-ji-shi-li/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
Golang AST语法树使用教程及示例
许多自动化代码生成工具都离不开语法树分析，例如goimport，gomock，wire等项目都离不开语法树分析。基于语法树分析，可以实现许多有趣实用的工具。本篇将结合示例，展示如何基于ast标准包操作语法树。
本篇中的代码的完整示例可以在这里找到：ast-example
首先我们看下语法树长什么样子，以下代码将打印./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：
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参数
实现这个功能我们需要四步：
- 遍历整个语法树
- 判断是否已经import
context包，如果没有则import - 遍历所有的接口方法，判断方法列表中是否有
context.Context类型的入参，如果没有我们将其添加到方法的第一个参数 - 将修改过后的语法树转换成Go代码并输出
遍历语法树
语法树层级较深，嵌套关系复杂，如果不能完全掌握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等等)。
syslog.ravelin.com/how-to-make…
medium.com/@astrid.deg…
stackoverflow.com/questions/3…
github.com/golang/go/i…
github.com/golang/go/i…
golang.org/src/go/ast/…
Recommend
About Joyk
Aggregate valuable and interesting links.
Joyk means Joy of geeK