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

ast

本篇中的代码的完整示例能够在这里找到：ast-examplegit

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:golang

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

首先是文件所属的包名，和其声明在文件中的位置：bash

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

...
     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  .  .  }
 ....
复制代码

Decl*ast.GenDeclimport

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()
	}
....
)
复制代码

NodeDeclsNode$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
	}
...
复制代码

IndentName

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
	}
...
复制代码

StructTypeFuncTypeInterfaceType

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
	}
...
复制代码

mainAssignStmt

...
 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

ImportSpecValueSpecTypeSpec

$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
	}
复制代码

ImportSpecValueSpecTypeSpecImportSpecTypeSpec

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
	}
)
...
复制代码

BadDeclGenDeclFunDeclGenDeclFunDecl

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参数

实现这个功能咱们须要四步：

contextcontext.Context

遍历语法树

ast

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

func Walk(v Visitor, node Node) 
复制代码

type Visitor interface {
	Visit(node Node) (w Visitor)
}
复制代码

WalkVisitorWalkVisitor.VisitVisitnilVisitor

// 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代码

formatformat.Nodegofmt

...
	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)
}
复制代码

astPos()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中有具体的讨论，官方认可这是一个设计缺陷，但仍是迟迟未能改进。其中有位火烧眉毛的小哥提供了本身的方案：

若是实在是要对有注释的语法树进行修改，能够尝试一下。 虽然语法树的确存在修改困难问题，但其仍是能知足大部分基于语法树分析的代码生成工做了(gomock,wire等等)。

参考