Understanding Go AST Through Automated Generation with Stringer

golang 18 minutes read (About 2680 words)

Unleash the Potential of Your Go Code with go ast and Automation

Contact Author@：Medium,LinkedIn,Twitter

Go AST (from github.com/MariaLetta/free-gophers-pack)|300

Hello, here is Wesley, Today’s article is about ast in Go. Without further ado, let’s get started.💪

Recently, I learned about using stringer and stumbled upon related content for golang ast, which led me to discover that I can also write some Go-related automation tools.

1.1 stringer and go generate

In development, we often need to define error codes with unique identifiers for specific error information. Additionally, we need to set the description for each error code. For example, in HTTP protocol, 200 represents “OK”, while 404 represents “Not Found”.

Traditionally, when defining error codes, we would also need to add a description, which is often forgotten and makes maintenance difficult. However, we can elegantly solve this problem using go generate + stringer toolchain.

stringer usage:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
Usage of stringer:
        stringer [flags] -type T [directory]
        stringer [flags] -type T files...  # Must be a single package
For more information, see:
        https://pkg.go.dev/golang.org/x/tools/cmd/stringer
Flags:
   -linecomment
        use line comment text as printed text when present
   -output string
        output file name; default srcdir/<type>_string.go
   -tags string
        comma-separated list of build tags to apply
   -trimprefix prefix
        trim the prefix from the generated constant names
   -type string
        comma-separated list of type names; must be set

Example:

1
2
3
4
5
6
7
8
9
package main

//go:generate stringer -type ErrCode -trimprefix ERR_CODE_ -output code_string.go
type ErrCode int
const (
        ERR_CODE_OK             ErrCode = 0 // ok
        ERR_CODE_INVALID_PARAMS ErrCode = 1 // invalid params
        ERR_CODE_TIMEOUT        ErrCode = 2 // timeout
)

Running go generate generates the following code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
// Code generated by "stringer -type ErrCode -trimprefix ERR_CODE -output code_string.go"; DO NOT EDIT.

package main

import "strconv"

func _() {
        // An "invalid array index" compiler error signifies that the constant values have changed.
        // Re-run the stringer command to generate them again.
        var x [1]struct{}
        _ = x[ERR_CODE_OK-0]
        _ = x[ERR_CODE_INVALID_PARAMS-1]
        _ = x[ERR_CODE_TIMEOUT-2]
}

const _ErrCode_name = "_OK_INVALID_PARAMS_TIMEOUT"

var _ErrCode_index = [...]uint8{0, 3, 18, 26}

func (i ErrCode) String() string {
        if i < 0 || i >= ErrCode(len(_ErrCode_index)-1) {
                return "ErrCode(" + strconv.FormatInt(int64(i), 10) + ")"
        }
        return _ErrCode_name[_ErrCode_index[i]:_ErrCode_index[i+1]]
}

The go generate command can be used after reading the article Generating code - The Go Programming Language and running go help generate to see basic usage.

1.2 How does stringer implement automatic generation?

Source code implementation: stringer.go

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/constant"
	"go/format"
	"go/token"
	"go/types"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"golang.org/x/tools/go/packages"
)

The main parts are “flag” and “go/ast”, “go/token”. The “flag” package mainly handles some command-line arguments, which can be referred to in my article: Flag Package You should Know in Golang | by Wesley Wei | Programmer’s Career.

The key point is to understand “go/ast”, which is a powerful tool for processing source code in the Go language. It not only simplifies the process of analyzing code, but also provides strong support for code generation and reconstruction.

1.3 Abstract Syntax Tree (AST)

Those who have studied compiler theory may be familiar with the front-end and back-end of a compiler. The front-end typically handles lexical analysis, syntax analysis, type checking, and intermediate code generation, while the back-end mainly focuses on generating and optimizing target code, which is translating intermediate code into binary machine code that can be executed by the target machine.

However, for those who are not familiar with compilers, we only need to understand the front-end part, which does not involve the back-end. Additionally, the official Go library go/ast provides a ready-to-use package.

1.3.1 Concept and Applications

An Abstract Syntax Tree (AST) is a tree-like data structure used to represent the syntactic structure of source code. AST breaks down the source code into nodes with hierarchical structures, where each node represents a syntactic element in the code, such as expressions, statements, variables, etc.

AST has wide applications in compilers and analysis tools:

  1. Code Analysis: The compiler converts the source code to an AST, providing a foundation for subsequent code analysis and optimization.
  2. Code Analysis and Checking: By traversing the AST, you can check the code for errors or improvement points using static analysis tools or code checking tools.
  3. Code Transformation and Optimization: By modifying the AST, you can implement automatic code optimization and transformation, such as code formatting tools.
  4. Code Generation: You can write code generation tools to generate new source code files or intermediate code.

1.3.2 How to Generate and Use AST in Go

In the Go language, we mainly use the following three standard libraries to generate and operate on AST:

  • go/parser: Used to parse Go source code and generate an AST.
  • go/token: Used to manage the positions and tokens of the source code.
  • go/ast: Used to represent and operate on the AST.

We can use the parser.ParseFile function to parse a file and generate an AST. Here is a simple example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
package main

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

func main() {
    src := `package main
    import "fmt"
    func main() {
        fmt.Println("Hello, World!")
    }`

    // Create a new file set
    fset := token.NewFileSet()

    // Parse the source code and generate an AST
    node, err := parser.ParseFile(fset, "", src, parser.ParseComments)
    if err != nil {
        fmt.Println("Error:", err)
        return
    }

    // Print the generated AST
    ast.Print(fset, node)
}

output:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
 0  *ast.File {
 1  .  Package: 1:1
 2  .  Name: *ast.Ident {
 3  .  .  NamePos: 1:9
 4  .  .  Name: "main"
 5  .  }
 6  .  Decls: []ast.Decl (len = 2) {
 7  .  .  0: *ast.GenDecl {
 8  .  .  .  TokPos: 3:1
 9  .  .  .  Tok: import
10  .  .  .  Lparen: -
11  .  .  .  Specs: []ast.Spec (len = 1) {
12  .  .  .  .  0: *ast.ImportSpec {
13  .  .  .  .  .  Path: *ast.BasicLit {
14  .  .  .  .  .  .  ValuePos: 3:8
15  .  .  .  .  .  .  Kind: STRING
16  .  .  .  .  .  .  Value: "\"fmt\""
17  .  .  .  .  .  }
18  .  .  .  .  .  EndPos: -
19  .  .  .  .  }
20  .  .  .  }
21  .  .  .  Rparen: -
22  .  .  }
23  .  .  1: *ast.FuncDecl {
24  .  .  .  Name: *ast.Ident {
25  .  .  .  .  NamePos: 5:6
26  .  .  .  .  Name: "main"
27  .  .  .  .  Obj: *ast.Object {
28  .  .  .  .  .  Kind: func
29  .  .  .  .  .  Name: "main"
30  .  .  .  .  .  Decl: *(obj @ 23)
31  .  .  .  .  }
32  .  .  .  }
33  .  .  .  Type: *ast.FuncType {
34  .  .  .  .  Func: 5:1
35  .  .  .  .  Params: *ast.FieldList {
36  .  .  .  .  .  Opening: 5:10
37  .  .  .  .  .  Closing: 5:11
38  .  .  .  .  }
39  .  .  .  }
40  .  .  .  Body: *ast.BlockStmt {
41  .  .  .  .  Lbrace: 5:13
42  .  .  .  .  List: []ast.Stmt (len = 1) {
43  .  .  .  .  .  0: *ast.ExprStmt {
44  .  .  .  .  .  .  X: *ast.CallExpr {
45  .  .  .  .  .  .  .  Fun: *ast.SelectorExpr {
46  .  .  .  .  .  .  .  .  X: *ast.Ident {
47  .  .  .  .  .  .  .  .  .  NamePos: 6:5
48  .  .  .  .  .  .  .  .  .  Name: "fmt"
49  .  .  .  .  .  .  .  .  }
50  .  .  .  .  .  .  .  .  Sel: *ast.Ident {
51  .  .  .  .  .  .  .  .  .  NamePos: 6:9
52  .  .  .  .  .  .  .  .  .  Name: "Println"
53  .  .  .  .  .  .  .  .  }
54  .  .  .  .  .  .  .  }
55  .  .  .  .  .  .  .  Lparen: 6:16
56  .  .  .  .  .  .  .  Args: []ast.Expr (len = 1) {
57  .  .  .  .  .  .  .  .  0: *ast.BasicLit {
58  .  .  .  .  .  .  .  .  .  ValuePos: 6:17
59  .  .  .  .  .  .  .  .  .  Kind: STRING
60  .  .  .  .  .  .  .  .  .  Value: "\"Hello, World!\""
61  .  .  .  .  .  .  .  .  }
62  .  .  .  .  .  .  .  }
63  .  .  .  .  .  .  .  Ellipsis: -
64  .  .  .  .  .  .  .  Rparen: 6:32
65  .  .  .  .  .  .  }
66  .  .  .  .  .  }
67  .  .  .  .  }
68  .  .  .  .  Rbrace: 7:1
69  .  .  .  }
70  .  .  }
71  .  }
72  .  FileStart: 1:1
73  .  FileEnd: 7:2
74  .  Scope: *ast.Scope {
75  .  .  Objects: map[string]*ast.Object (len = 1) {
76  .  .  .  "main": *(obj @ 27)
77  .  .  }
78  .  }
79  .  Imports: []*ast.ImportSpec (len = 1) {
80  .  .  0: *(obj @ 12)
81  .  }
82  .  Unresolved: []*ast.Ident (len = 1) {
83  .  .  0: *(obj @ 46)
84  .  }
85  .  GoVersion: ""
86  }

Common Node Types

  1. ast.File: Represents a source file.
  2. ast.GenDecl: Represents general declarations (such as imports, constants, variables, type declarations).
  3. ast.FuncDecl: Represents function declarations.
  4. ast.BlockStmt: Represents code blocks.
  5. ast.Expr: The base node type for all expressions, such as literals and binary expressions.
  6. ast.Stmt: The base node type for all statements, such as assignment statements and return statements.

For example, the *ast.FuncDecl output represents a function definition with the name main. The content is quite extensive, so you can refer to ast.go; l=972 for more information.

1.3.3 Traversing and Visiting AST Nodes

After generating the AST, you can traverse and visit each node. For example, the following code shows how to traverse the AST and extract the desired content:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
package main
import (
    "go/ast"
    "go/parser"
    "go/token"
    "fmt"
)

func main() {
    src := `package main
    func add(a int, b int) int {
        return a + b
    }`

    fset := token.NewFileSet()
    fileNode, err := parser.ParseFile(fset, "", src, parser.ParseComments)
    if err != nil {
        fmt.Println("Error:", err)
        return
    }

    ast.Inspect(fileNode, func(n ast.Node) bool {
        switch n := n.(type) {
        case *ast.File:
            fmt.Printf("Package: %s\n", n.Name.Name)
        case *ast.FuncDecl:
            fmt.Printf("Function: %s\n", n.Name.Name)
        case *ast.GenDecl:
            fmt.Printf("General Declaration: %s\n", n.Tok)
        }
        return true
    })
}

Output:

1
2
Package: main
Function: add

1.4 Actual Applications

1.4.1 Code Inspection

Code inspection is an essential part of static analysis, which can help detect potential issues or optimize code. By traversing the Abstract Syntax Tree (AST), we can write custom rules to inspect code.

The following example demonstrates how to use AST to check for unused variables:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
package main

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

func main() {
    src := `package main
    func main() {
        var unused int
        fmt.Println("Hello, World!")
    }`

    fset := token.NewFileSet()
    node, err := parser.ParseFile(fset, "", src, parser.ParseComments)
    if err != nil {
        fmt.Println("Error:", err)
        return
    }

    // Check for unused variables
    ast.Inspect(node, func(n ast.Node) bool {
        if decl, ok := n.(*ast.GenDecl); ok {
            for _, spec := range decl.Specs {
                if valueSpec, ok := spec.(*ast.ValueSpec); ok {
                    for _, name := range valueSpec.Names {
                        if !name.IsExported() && name.Name != "_" {
                            fmt.Printf("Unused variable: %s\n", name.Name)
                        }
                    }
                }
            }
        }
        return true
    })
}

1.4.2 Code Formatting

Code formatting is an important part of ensuring code style consistency. By using AST, we can reorganize the code to conform to predetermined format rules.

The following example demonstrates how to use AST to add comments for each function:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
package main
import (
    "fmt"
    "go/ast"
    "go/parser"
    "go/printer"
    "go/token"
    "os"
)
func main() {
    src := `package main
func add(a int, b int) int {
    return a + b
}`

    fset := token.NewFileSet()
    node, err := parser.ParseFile(fset, "", src, parser.ParseComments)
    if err != nil {
        fmt.Println("Error:", err)
        return
    }

    // Add comments for each function
    ast.Inspect(node, func(n ast.Node) bool {
        if decl, ok := n.(*ast.FuncDecl); ok {
            decl.Doc = &ast.CommentGroup{
                List: []*ast.Comment{
                    {
                        Slash: decl.Pos() - 1,
                        Text:   "// Function " + decl.Name.Name,
                    },
                },
            }
        }
        return true
    })

    printer.Fprint(os.Stdout, fset, node)
}

1.4.3 Comment Handling

Comment handling can be used for tasks such as extracting documentation and generating code. By traversing AST, we can collect and process comments in the code.

The following example demonstrates how to extract and print all comments:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
package main

import (  
    "fmt"  
    "go/ast"  
    "go/parser"  
    "go/token"  
)  
func main() {  
    src := `package main

// This is a package comment  
func add(a int, b int) int {  
     // This is a function comment  
    return a + b  
}`

    fset := token.NewFileSet()
    node, err := parser.ParseFile(fset, "", src, parser.ParseComments)
    if err != nil {
        fmt.Println("Error:", err)
        return
    }

    // Extract and print comments
    for _, comment := range node.Comments {
        fmt.Println(comment.Text())
    }

}

1.4.4 Concrete Examples in Go

You must have thought of some Go tools after seeing the previous examples.

In Go, many tools are implemented based on the go/ast (Abstract Syntax Tree) library, which is used to parse and analyze Go source code. Here are a few common tools that use go/ast:

gofmt

  • gofmt is a built-in code formatting tool for Go, which formats Go code into a standard style. It uses go/ast to analyze the code structure and convert it into a uniform format.
  • For example, gofmt automatically performs indentation, alignment, and whitespace operations to ensure consistency and readability of Go code.

go vet

  • go vet is a static code analysis tool that can detect potential errors or non-compliant usage in code. By analyzing the code using go/ast, go vet can detect common errors such as mismatched Printf parameters, incorrect lock operations, and more.

golint

  • golint is a code linting tool for Go that helps developers follow Go’s coding conventions. It uses go/ast to analyze the code and find non-compliant coding patterns, such as unnamed exported structures or naming conventions not followed.

go doc

  • go doc is a documentation tool for Go that uses go/ast to analyze code comments and generate function, method, type, and other documentation information.
  • This makes it possible for go doc to provide convenient documentation viewing options, especially for local code libraries.

gorename

  • gorename is a renaming tool that allows developers to safely rename Go code identifiers. It uses go/ast to analyze the code structure and dependencies, ensuring all references can be safely renamed.

stringer

  • stringer is an official Go code generation tool specifically designed for generating string representations of enumeration types. It uses go/ast to parse constant definitions in code and generate related string methods, making it easier to debug output and log recording in code.

These tools, based on go/ast, provide structured analysis capabilities for Go code, enabling functions such as formatting, static checking, documentation generation, and code generation.

1.5 Other Applications of AST Tools

We will introduce two projects, pingcap’s failpoint and Uber-go’s gopatch.

The usage of failpoint is to implement the Marker function in code. These functions are optimized away by the compiler at normal compilation time, so there is no extra overhead during normal runtime:

1
2
3
4
var outerVar = "declare in outer scope"
failpoint.Inject("failpoint-name-for-demo", func(val failpoint.Value) {
    fmt.Println("unit-test", val, outerVar)
})

When fault injection occurs, the Marker function is converted to a fault injection function using go ast.

Uber-go’s gopatch is also very powerful. For example, if your code has many go func-opened goroutines and you want to batch in recover logic, it can be tedious to do this manually. At this time, you can use gopatcher:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
var patchTemplateString = `@@  
@@  
+ import "runtime/debug"  
+ import "{{ Logger }}"  
+ import "{{ Statsd }}"  
   go func(...) {  
+    defer func() {  
+        if err := recover(); err != nil {  
+            statsd.Count1("{{ StatsdTag }}", "{{ FileName }}")  
+            logging.Error("{{ LoggerTag }}", "{{ FileName }} recover from panic, err=%+v, stack=%v", err, string(debug.Stack()))  
+        }  
+    }()  
    ...  
  }()`

You can write a template like this:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
package main

import (
    "os"
    "text/template"
)

type PatchTemplate struct {
    Statsd      string
    Logger      string
    StatsdTag   string
    FileName    string
    LoggerTag   string
}

func main() {
    patchTemplate := `@@
@@
+ import "runtime/debug"
+ import "{{ .Logger }}"
+ import "{{ .Statsd }}"

go func(...) {
+    defer func() {
+        if err := recover(); err != nil {
+            statsd.Count1("{{ .StatsdTag }}", "{{ .FileName }}")
+            logging.Error("{{ .LoggerTag }}", "{{ .FileName }} recover from panic, err=%+v, stack=%v", err, string(debug.Stack()))
+        }
+    }()

    ...  
  }()`

    data := PatchTemplate{
        Statsd:     "github.com/yourusername/statsd",
        Logger:     "github.com/yourusername/logger",
        StatsdTag:  "error_count",
        FileName:   "example.go",
        LoggerTag:  "example_logger",
    }

    t := template.Must(template.New("patch").Parse(patchTemplate))
    t.Execute(os.Stdout, data)
}

This example can automatically inject the recover statement block at the beginning of go func(…) {, which is very convenient.

Of course, gopatch can do many things, and you can explore it freely.

1.6 Conclusion

The above content introduced the basic usage of go ast and some mature tools, but in actual work, you will encounter scenarios that are between these two extremes, such as custom error codes, standard HTTP error codes being too few for a large-scale service, so you can write an automated tool to generate custom error codes.

Of course, automated tools are only extremely useful when dealing with a large amount of redundant work, as they can help us save labor and maintain focus.

In summary, the go ast library provides developers with a lot of space to handle Go code, mastering AST usage can enable more efficient code analysis and tool development.

More Series Articles about You Should Know In Golang:

https://wesley-wei.medium.com/list/you-should-know-in-golang-e9491363cd9a

And I’m Wesley, delighted to share knowledge from the world of programming. 

Don’t forget to follow me for more informative content, or feel free to share this with others who may also find it beneficial. it would be a great help to me.

Give me some free applauds, highlights, or replies, and I’ll pay attention to those reactions, which will determine whether or not I continue to post this type of article.

See you in the next article. 👋

中文文章: https://programmerscareer.com/zh-cn/golang-ast/
Author: Medium,LinkedIn,Twitter
Note: Originally written at https://programmerscareer.com/golang-ast/ at 2024-10-26 20:28.
Copyright: BY-NC-ND 3.0

wechat
Wechat(微信公众号)
x
#code
Enlightening Weekly 001: Understanding Code and Investment Insights SyncOnce You Should Know in Golang

Comments

Your browser is out-of-date!

Update your browser to view this website correctly. Update my browser now

×