Golang AST语法树使用教程及示例

https://juejin.im/post/5db7195df265da4d104b7fb7

许多自动化代码生成工具都离不开语法树分析,例如goimportgomockwire等项目都离不开语法树分析。基于语法树分析,可以实现许多有趣实用的工具。本篇将结合示例,展示如何基于ast标准包操作语法树。

本篇中的代码的完整示例可以在这里找到:ast-example

Quick Start

首先我们看下语法树长什么样子,以下代码将打印./demo.go文件的语法树:

  1. package main
  2. import (
  3. "go/ast"
  4. "go/parser"
  5. "go/token"
  6. "log"
  7. "path/filepath"
  8. )
  9. func main() {
  10. fset := token.NewFileSet()
  11. // 这里取绝对路径,方便打印出来的语法树可以转跳到编辑器
  12. path, _ := filepath.Abs("./demo.go")
  13. f, err := parser.ParseFile(fset, path, nil, parser.AllErrors)
  14. if err != nil {
  15. log.Println(err)
  16. return
  17. }
  18. // 打印语法树
  19. ast.Print(fset, f)
  20. }

demo.go:

  1. package main
  2. import (
  3. "context"
  4. )
  5. // Foo 结构体
  6. type Foo struct {
  7. i int
  8. }
  9. // Bar 接口
  10. type Bar interface {
  11. Do(ctx context.Context) error
  12. }
  13. // main方法
  14. func main() {
  15. a := 1
  16. }

demo.go文件已尽量简化,但其语法树的输出内容依旧十分庞大。我们截取部分来做一些简要的说明。

首先是文件所属的包名,和其声明在文件中的位置:

  1. 0 *ast.File {
  2. 1 . Package: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:1
  3. 2 . Name: *ast.Ident {
  4. 3 . . NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:9
  5. 4 . . Name: "main"
  6. 5 . }
  7. ...

紧接着是Decls,也就是Declarations,其包含了声明的一些变量,方法,接口等:

  1. ...
  2. 6 . Decls: []ast.Decl (len = 4) {
  3. 7 . . 0: *ast.GenDecl {
  4. 8 . . . TokPos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:3:1
  5. 9 . . . Tok: import
  6. 10 . . . Lparen: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:3:8
  7. 11 . . . Specs: []ast.Spec (len = 1) {
  8. 12 . . . . 0: *ast.ImportSpec {
  9. 13 . . . . . Path: *ast.BasicLit {
  10. 14 . . . . . . ValuePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:4:2
  11. 15 . . . . . . Kind: STRING
  12. 16 . . . . . . Value: "\"context\""
  13. 17 . . . . . }
  14. 18 . . . . . EndPos: -
  15. 19 . . . . }
  16. 20 . . . }
  17. 21 . . . Rparen: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:5:1
  18. 22 . . }
  19. ....

可以看到该语法树包含了4条Decl记录,我们取第一条记录为例,该记录为*ast.GenDecl类型。不难看出这条记录对应的是我们的import代码段。始位置(TokPos),左右括号的位置(Lparen,Rparen),和import的包(Specs)等信息都能从语法树中得到。

语法树的打印信来自ast.File结构体:

$GOROOT/src/go/ast/ast.go

  1. // 该结构体位于标准包 go/ast/ast.go 中,有兴趣可以转跳到源码阅读更详尽的注释
  2. type File struct {
  3. Doc *CommentGroup // associated documentation; or nil
  4. Package token.Pos // position of "package" keyword
  5. Name *Ident // package name
  6. Decls []Decl // top-level declarations; or nil
  7. Scope *Scope // package scope (this file only)
  8. Imports []*ImportSpec // imports in this file
  9. Unresolved []*Ident // unresolved identifiers in this file
  10. Comments []*CommentGroup // list of all comments in the source file
  11. }

结合注释和字段名我们大概知道每个字段的含义,接下来我们详细梳理一下语法树的组成结构。

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

  1. ...
  2. // All node types implement the Node interface.
  3. type Node interface {
  4. Pos() token.Pos // position of first character belonging to the node
  5. End() token.Pos // position of first character immediately after the node
  6. }
  7. // All expression nodes implement the Expr interface.
  8. type Expr interface {
  9. Node
  10. exprNode()
  11. }
  12. // All statement nodes implement the Stmt interface.
  13. type Stmt interface {
  14. Node
  15. stmtNode()
  16. }
  17. // All declaration nodes implement the Decl interface.
  18. type Decl interface {
  19. Node
  20. declNode()
  21. }
  22. ...
  23. // A Spec node represents a single (non-parenthesized) import,
  24. // constant, type, or variable declaration.
  25. //
  26. type (
  27. // The Spec type stands for any of *ImportSpec, *ValueSpec, and *TypeSpec.
  28. Spec interface {
  29. Node
  30. specNode()
  31. }
  32. ....
  33. )

可以看到所有的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

  1. ...
  2. // An Ident node represents an identifier.
  3. Ident struct {
  4. NamePos token.Pos // identifier position
  5. Name string // identifier name
  6. Obj *Object // denoted object; or nil
  7. }
  8. ...

Indent(identifier)表示一个标识符,比如Quick Start示例中表示包名的Name字段就是一个expression node:

  1. 0 *ast.File {
  2. 1 . Package: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:1
  3. 2 . Name: *ast.Ident { <----
  4. 3 . . NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:1:9
  5. 4 . . Name: "main"
  6. 5 . }
  7. ...

接下来是type node。

$GOROOT/src/go/ast/ast.go

  1. ...
  2. // A StructType node represents a struct type.
  3. StructType struct {
  4. Struct token.Pos // position of "struct" keyword
  5. Fields *FieldList // list of field declarations
  6. Incomplete bool // true if (source) fields are missing in the Fields list
  7. }
  8. // Pointer types are represented via StarExpr nodes.
  9. // A FuncType node represents a function type.
  10. FuncType struct {
  11. Func token.Pos // position of "func" keyword (token.NoPos if there is no "func")
  12. Params *FieldList // (incoming) parameters; non-nil
  13. Results *FieldList // (outgoing) results; or nil
  14. }
  15. // An InterfaceType node represents an interface type.
  16. InterfaceType struct {
  17. Interface token.Pos // position of "interface" keyword
  18. Methods *FieldList // list of methods
  19. Incomplete bool // true if (source) methods are missing in the Methods list
  20. }
  21. ...

type node很好理解,它包含一些复合类型,例如在Quick Start中出现的StructType,FuncTypeInterfaceType

Statement

赋值语句,控制语句(if,else,for,select…)等均属于statement node。

$GOROOT/src/go/ast/ast.go

  1. ...
  2. // An AssignStmt node represents an assignment or
  3. // a short variable declaration.
  4. //
  5. AssignStmt struct {
  6. Lhs []Expr
  7. TokPos token.Pos // position of Tok
  8. Tok token.Token // assignment token, DEFINE
  9. Rhs []Expr
  10. }
  11. ...
  12. // An IfStmt node represents an if statement.
  13. IfStmt struct {
  14. If token.Pos // position of "if" keyword
  15. Init Stmt // initialization statement; or nil
  16. Cond Expr // condition
  17. Body *BlockStmt
  18. Else Stmt // else branch; or nil
  19. }
  20. ...

例如Quick Start中,我们在main函数中对变量a赋值的程序片段就属于AssignStmt:

  1. ...
  2. 174 . . . Body: *ast.BlockStmt {
  3. 175 . . . . Lbrace: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:18:13
  4. 176 . . . . List: []ast.Stmt (len = 1) {
  5. 177 . . . . . 0: *ast.AssignStmt { <--- 这里
  6. 178 . . . . . . Lhs: []ast.Expr (len = 1) {
  7. 179 . . . . . . . 0: *ast.Ident {
  8. 180 . . . . . . . . NamePos: /usr/local/gopath/src/github.com/DrmagicE/ast-example/quickstart/demo.go:19:2
  9. 181 . . . . . . . . Name: "a"
  10. ...

Spec Node

Spec node只有3种,分别是ImportSpecValueSpecTypeSpec

$GOROOT/src/go/ast/ast.go

  1. // An ImportSpec node represents a single package import.
  2. ImportSpec struct {
  3. Doc *CommentGroup // associated documentation; or nil
  4. Name *Ident // local package name (including "."); or nil
  5. Path *BasicLit // import path
  6. Comment *CommentGroup // line comments; or nil
  7. EndPos token.Pos // end of spec (overrides Path.Pos if nonzero)
  8. }
  9. // A ValueSpec node represents a constant or variable declaration
  10. // (ConstSpec or VarSpec production).
  11. //
  12. ValueSpec struct {
  13. Doc *CommentGroup // associated documentation; or nil
  14. Names []*Ident // value names (len(Names) > 0)
  15. Type Expr // value type; or nil
  16. Values []Expr // initial values; or nil
  17. Comment *CommentGroup // line comments; or nil
  18. }
  19. // A TypeSpec node represents a type declaration (TypeSpec production).
  20. TypeSpec struct {
  21. Doc *CommentGroup // associated documentation; or nil
  22. Name *Ident // type name
  23. Assign token.Pos // position of '=', if any
  24. Type Expr // *Ident, *ParenExpr, *SelectorExpr, *StarExpr, or any of the *XxxTypes
  25. Comment *CommentGroup // line comments; or nil
  26. }

ImportSpec表示一个单独的import,ValueSpec表示一个常量或变量的声明,TypeSpec则表示一个type声明。例如 在Quick Start示例中,出现了ImportSpecTypeSpec

  1. import (
  2. "context" // <--- 这里是一个ImportSpec node
  3. )
  4. // Foo 结构体
  5. type Foo struct { // <--- 这里是一个TypeSpec node
  6. i int
  7. }

在语法树的打印结果中可以看到对应的输出,小伙伴们可自行查找。

Declaration Node

Declaration node也只有三种:

$GOROOT/src/go/ast/ast.go

  1. ...
  2. type (
  3. // A BadDecl node is a placeholder for declarations containing
  4. // syntax errors for which no correct declaration nodes can be
  5. // created.
  6. //
  7. BadDecl struct {
  8. From, To token.Pos // position range of bad declaration
  9. }
  10. // A GenDecl node (generic declaration node) represents an import,
  11. // constant, type or variable declaration. A valid Lparen position
  12. // (Lparen.IsValid()) indicates a parenthesized declaration.
  13. //
  14. // Relationship between Tok value and Specs element type:
  15. //
  16. // token.IMPORT *ImportSpec
  17. // token.CONST *ValueSpec
  18. // token.TYPE *TypeSpec
  19. // token.VAR *ValueSpec
  20. //
  21. GenDecl struct {
  22. Doc *CommentGroup // associated documentation; or nil
  23. TokPos token.Pos // position of Tok
  24. Tok token.Token // IMPORT, CONST, TYPE, VAR
  25. Lparen token.Pos // position of '(', if any
  26. Specs []Spec
  27. Rparen token.Pos // position of ')', if any
  28. }
  29. // A FuncDecl node represents a function declaration.
  30. FuncDecl struct {
  31. Doc *CommentGroup // associated documentation; or nil
  32. Recv *FieldList // receiver (methods); or nil (functions)
  33. Name *Ident // function/method name
  34. Type *FuncType // function signature: parameters, results, and position of "func" keyword
  35. Body *BlockStmt // function body; or nil for external (non-Go) function
  36. }
  37. )
  38. ...

BadDecl表示一个有语法错误的节点; GenDecl用于表示import, const,type或变量声明;FunDecl用于表示函数声明。 GenDeclFunDecl在Quick Start例子中均有出现,小伙伴们可自行查找。

Common Node

除去上述四种类别划分的node,还有一些node不属于上面四种类别:

$GOROOT/src/go/ast/ast.go

  1. // Comment 注释节点,代表单行的 //-格式 或 /*-格式的注释.
  2. type Comment struct {
  3. ...
  4. }
  5. ...
  6. // CommentGroup 注释块节点,包含多个连续的Comment
  7. type CommentGroup struct {
  8. ...
  9. }
  10. // Field 字段节点, 可以代表结构体定义中的字段,接口定义中的方法列表,函数前面中的入参和返回值字段
  11. type Field struct {
  12. ...
  13. }
  14. ...
  15. // FieldList 包含多个Field
  16. type FieldList struct {
  17. ...
  18. }
  19. // File 表示一个文件节点
  20. type File struct {
  21. ...
  22. }
  23. // Package 表示一个包节点
  24. type Package struct {
  25. ...
  26. }

Quick Start示例包含了上面列举的所有node,小伙伴们可以自行查找。更为详细的注释和具体的结构体字段请查阅源码。

所有的节点类型大致列举完毕,其中还有许多具体的节点类型未能一一列举,但基本上都是大同小异,源码注释也比较清晰,等用到的时候再细看也不迟。现在我们对整个语法树的构造有了基本的了解,接下来通过几个示例来演示具体用法。

示例

为文件中所有接口方法添加context参数

实现这个功能我们需要四步:

  1. 遍历整个语法树
  2. 判断是否已经importcontext包,如果没有则import
  3. 遍历所有的接口方法,判断方法列表中是否有context.Context类型的入参,如果没有我们将其添加到方法的第一个参数
  4. 将修改过后的语法树转换成Go代码并输出
遍历语法树

语法树层级较深,嵌套关系复杂,如果不能完全掌握node之间的关系和嵌套规则,我们很难自己写出正确的遍历方法。不过好在ast包已经为我们提供了遍历方法:

$GOROOT/src/go/ast/ast.go

  1. func Walk(v Visitor, node Node)
  1. type Visitor interface {
  2. Visit(node Node) (w Visitor)
  3. }

Walk方法会按照深度优先搜索方法(depth-first order)遍历整个语法树,我们只需按照我们的业务需要,实现Visitor接口即可。 Walk每遍历一个节点就会调用Visitor.Visit方法,传入当前节点。如果Visit返回nil,则停止遍历当前节点的子节点。本示例的Visitor实现如下:

  1. // Visitor
  2. type Visitor struct {
  3. }
  4. func (v *Visitor) Visit(node ast.Node) ast.Visitor {
  5. switch node.(type) {
  6. case *ast.GenDecl:
  7. genDecl := node.(*ast.GenDecl)
  8. // 查找有没有import context包
  9. // Notice:没有考虑没有import任何包的情况
  10. if genDecl.Tok == token.IMPORT {
  11. v.addImport(genDecl)
  12. // 不需要再遍历子树
  13. return nil
  14. }
  15. case *ast.InterfaceType:
  16. // 遍历所有的接口类型
  17. iface := node.(*ast.InterfaceType)
  18. addContext(iface)
  19. // 不需要再遍历子树
  20. return nil
  21. }
  22. return v
  23. }
添加import
  1. // addImport 引入context包
  2. func (v *Visitor) addImport(genDecl *ast.GenDecl) {
  3. // 是否已经import
  4. hasImported := false
  5. for _, v := range genDecl.Specs {
  6. imptSpec := v.(*ast.ImportSpec)
  7. // 如果已经包含"context"
  8. if imptSpec.Path.Value == strconv.Quote("context") {
  9. hasImported = true
  10. }
  11. }
  12. // 如果没有import context,则import
  13. if !hasImported {
  14. genDecl.Specs = append(genDecl.Specs, &ast.ImportSpec{
  15. Path: &ast.BasicLit{
  16. Kind: token.STRING,
  17. Value: strconv.Quote("context"),
  18. },
  19. })
  20. }
  21. }
为接口方法添加参数
  1. // addContext 添加context参数
  2. func addContext(iface *ast.InterfaceType) {
  3. // 接口方法不为空时,遍历接口方法
  4. if iface.Methods != nil || iface.Methods.List != nil {
  5. for _, v := range iface.Methods.List {
  6. ft := v.Type.(*ast.FuncType)
  7. hasContext := false
  8. // 判断参数中是否包含context.Context类型
  9. for _, v := range ft.Params.List {
  10. if expr, ok := v.Type.(*ast.SelectorExpr); ok {
  11. if ident, ok := expr.X.(*ast.Ident); ok {
  12. if ident.Name == "context" {
  13. hasContext = true
  14. }
  15. }
  16. }
  17. }
  18. // 为没有context参数的方法添加context参数
  19. if !hasContext {
  20. ctxField := &ast.Field{
  21. Names: []*ast.Ident{
  22. ast.NewIdent("ctx"),
  23. },
  24. // Notice: 没有考虑import别名的情况
  25. Type: &ast.SelectorExpr{
  26. X: ast.NewIdent("context"),
  27. Sel: ast.NewIdent("Context"),
  28. },
  29. }
  30. list := []*ast.Field{
  31. ctxField,
  32. }
  33. ft.Params.List = append(list, ft.Params.List...)
  34. }
  35. }
  36. }
  37. }
将语法树转换成Go代码

format包为我们提供了转换函数,format.Node会将语法树按照gofmt的格式输出:

  1. ...
  2. var output []byte
  3. buffer := bytes.NewBuffer(output)
  4. err = format.Node(buffer, fset, f)
  5. if err != nil {
  6. log.Fatal(err)
  7. }
  8. // 输出Go代码
  9. fmt.Println(buffer.String())
  10. ...

输出结果如下:

  1. package main
  2. import (
  3. "context"
  4. )
  5. type Foo interface {
  6. FooA(ctx context.Context, i int)
  7. FooB(ctx context.Context, j int)
  8. FooC(ctx context.Context)
  9. }
  10. type Bar interface {
  11. BarA(ctx context.Context, i int)
  12. BarB(ctx context.Context)
  13. BarC(ctx context.Context)
  14. }

可以看到我们所有的接口方的第一个参数都变成了context.Context。建议将示例中的语法树先打印出来,再对照着代码看,方便理解。

一些坑与不足

至此我们已经完成了语法树的解析,遍历,修改以及输出。但细心的小伙伴可能已经发现:示例中的文件并没有出现一行注释。这的确是有意为之,如果我们加上注释,会发现最终生成文件的注释就像迷途的羔羊,完全找不到自己的位置。比如这样:

  1. //修改前
  2. type Foo interface {
  3. FooA(i int)
  4. // FooB
  5. FooB(j int)
  6. FooC(ctx context.Context)
  7. }
  8. // 修改后
  9. type Foo interface {
  10. FooA(ctx context.
  11. // FooB
  12. Context, i int)
  13. FooB(ctx context.Context, j int)
  14. FooC(ctx context.Context)
  15. }

导致这种现象的原因在于:ast包生成的语法树中的注释是”free-floating”的。还记得每个node都有Pos()End()方法来标识其位置吗?对于非注释节点,语法树能够正确的调整他们的位置,但却不能自动调整注释节点的位置。如果我们想要让注释出现在正确的位置上,我们必须手动设置节点PosEnd。源码注释中提到了这个问题:

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中有具体的讨论,官方承认这是一个设计缺陷,但还是迟迟未能改进。其中有位迫不及待的小哥提供了自己的方案:

github.com/dave/dst

如果实在是要对有注释的语法树进行修改,可以尝试一下。 虽然语法树的确存在修改困难问题,但其还是能满足大部分基于语法树分析的代码生成工作了(gomock,wire等等)。

参考

ft_authoradmin  ft_create_time2019-11-19 14:00
 ft_update_time2019-11-19 14:22