go generate and ast

楔(xiē)子

最近写API CURD比较多,为了结构清晰,返回值需要统一错误码,所以在一个统一的errcode包中定义错误码常量,以及其错误信息.

如下图所示,由于常量是导出字符 -> golint 检测需要编写注释 -> 注释信息其实就是错误信息,已经在下文的msg map[int]string中定义,如果在写就得写两遍

不写,就满屏波浪线,不能忍!

写了,就得Copy一份,还不利于维护,不能忍!

能不能只写一份注释,剩下的msg通过读取注释信息自动生成,将我们宝(hua)贵(diao)的生命,从这些重复繁杂无意义的劳动中解放出来。

为了实现这个伟大的目标, 需要以下两个关键的数据:

  1. 解析源代码获取常量与注释之间的关系 -> 🌲Go 抽象语法树: AST[3]
  2. 从 Go 源码生成 Go 代码 -> 👏 go generate[5]

👏 go generate

golang1.4版本中引入了go generate命令,常用于文件生成,例如在 Golang 官方博客[5]中介绍的Stringer可以为枚举自动实现Stringer的方法,从业务代码中解放出来

💻 命令文档

使用go help generate我们可以查看一下命令的帮助文档

1
2
3
▶ go help generate
usage: go generate [-run regexp] [-n] [-v] [-x] [build flags] [file.go... | packages]
...

解释很长,就不贴上来了,简要的概括一下:

  1. 参数说明

    • -run 正则表达式匹配命令行,仅执行匹配的命令(和go test -run类似)
    • -v 打印  已被检索处理的文件。
    • -n 打印出将被执行的命令,此时将不真实执行命令
    • -x 打印已执行的命令
  2. 举个栗子

    1
    2
    3
    4
    # 对当前包下的Go文件进行处理, 并打印已被检索处理的文件。
    go generate -v
    # 打印当前目录下所有文件中将要被执行的命令(实际不会执行)
    go generate -n ./...
  3. go generate会扫描.go源码文件中的注释//go:generate command args..., 并且执行其命令,注意:

    • 这些命令是为了更新或者创建 Go 源文件
    • command必须是可执行的指令,例如在 PATH 中或者使用绝对路径
    • arg如果带引号会被识别成一个参数, 例如: //go:generate command "x1 x2", 这条语句执行的命令只有一个参数
    • 注释中//go之间没有空格
  4. go generate必须手动执行,如果想等着go build, go test, go run 命令执行的时候自动执行,可以洗洗睡了

  5. 为了让别人或者是 IDE 识别代码是通过go generate生成的,请在生成的代码中添加注释(一般放在文件开头)

    1
    2
    # PS: 这是一个正则表达式
    ^// Code generated .* DO NOT EDIT\.$

    举个栗子:

    1
    2
    3
    // Code generated by mohuishou DO NOT EDIT

    package painkiller
  6. go generate在执行的时候会自动注入以下环境变量:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    $GOARCH
    系统架构: arm, amd64 等
    $GOOS
    操作系统: linux, windows 等
    $GOFILE
    当前执行的命令所处的文件名
    $GOLINE
    当前执行的命令在文件中的行号
    $GOPACKAGE
    执行的命令所处的文件的包名
    $DOLLAR
    $ 符号

🌰 Go 官方博客中给出的栗子

源文件: painkiller.go

1
2
3
4
5
6
7
8
9
10
11
12
13
//go:generate stringer -type=Pill

package painkiller

type Pill int

const (
Placebo Pill = iota
Aspirin
Ibuprofen
Paracetamol
Acetaminophen = Paracetamol
)

执行命令

1
go generate

生成文件: painkiller_stringer.go

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// generated by stringer -type Pill pill.go; DO NOT EDIT

package painkiller

import "fmt"

const _Pill_name = "PlaceboAspirinIbuprofenParacetamol"

var _Pill_index = [...]uint8{0, 7, 14, 23, 34}

func (i Pill) String() string {
if i < 0 || i+1 >= Pill(len(_Pill_index)) {
return fmt.Sprintf("Pill(%d)", i)
}
return _Pill_name[_Pill_index[i]:_Pill_index[i+1]]
}

从上面的 🌰,我们可以发现,在.go源文件中,添加了一行注释go:generate stringer -type=Pill, 执行命令go generate就调用stringer命令在同目录下生成了一个新的_stringer.go的文件

回想一下上文提到的需求,是不是感觉很类似,从 Go 源文件中,生成了一些不想重复写的业务逻辑

🌲 AST

回到前面的需求,我们需要从源代码中获取常量和注释之前的关系,这时就需要我们的 🌲AST 隆重登场了。

本文不对 AST 过多介绍,可以阅读参考资料中的 AST 标准库文档[3],Go 的 AST(抽象语法树)[4]

简要介绍一下 AST 包

基础的接口类型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
// Node AST树节点
type Node interface {
Pos() token.Pos
End() token.Pos
}

// Expr 所有的表达式都需要实现Expr接口
type Expr interface {
Node
exprNode()
}

// Stmt 所有的语句都需要实现Stmt接口
type Stmt interface {
Node
stmtNode()
}

// Decl 所有的声明都需要实现Decl接口
type Decl interface {
Node
declNode()
}

等会儿可能会用到的ValueSpec

1
2
3
4
5
6
7
8
// ValueSpec 表示常量声明或者变量声明
type 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
}

CommentMap

在 godoc[3]的 Example 中可以发现有一个CommentMap例子

1
2
// CommentMap把AST节点和其关联的注释列表进行映射
type CommentMap map[Node][]*CommentGroup
  1. 通过parse读取源码创建一个 AST

    1
    2
    3
    4
    5
    fset := token.NewFileSet() // positions are relative to fset
    f, err := parser.ParseFile(fset, "src.go", src, parser.ParseComments)
    if err != nil {
    panic(err)
    }
  2. 从 AST 中新建一个CommentMap

    1
    cmap := ast.NewCommentMap(fset, f, f.Comments)

需求实现

1. 获取常量和注释的关联关系

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
file := os.Getenv("GOFILE")
// 保存注释信息
var comments = make(map[string]string)

// 解析代码源文件,获取常量和注释之间的关系
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
checkErr(err)

// Create an ast.CommentMap from the ast.File's comments.
// This helps keeping the association between comments
// and AST nodes.
cmap := ast.NewCommentMap(fset, f, f.Comments)
for node := range cmap {
// 仅支持一条声明语句,一个常量的情况
if spec, ok := node.(*ast.ValueSpec); ok && len(spec.Names) == 1 {
// 仅提取常量的注释
ident := spec.Names[0]
if ident.Obj.Kind == ast.Con {
// 获取注释信息
comments[ident.Name] = getComment(ident.Name, spec.Doc)
}
}
}

2. 获取注释信息

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
// getComment 获取注释信息,来自AST标准库的summary方法
func getComment(name string, group *ast.CommentGroup) string {
var buf bytes.Buffer


for _, comment := range group.List {
// 注释信息会以 // 参数名,开始,我们实际使用时不需要,去掉
text := strings.TrimSpace(strings.TrimPrefix(comment.Text, fmt.Sprintf("// %s", name)))
buf.WriteString(text)
}

// replace any invisibles with blanks
bytes := buf.Bytes()
for i, b := range bytes {
switch b {
case '\t', '\n', '\r':
bytes[i] = ' '
}
}

return string(bytes)
}

3. 生成代码

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
const suffix = "_msg_gen.go"

// tpl 生成代码需要用到模板
const tpl = `
// Code generated by github.com/mohuishou/gen-const-msg DO NOT EDIT

// {{.pkg}} const code comment msg
package {{.pkg}}

// noErrorMsg if code is not found, GetMsg will return this
const noErrorMsg = "unknown error"

// messages get msg from const comment
var messages = map[int]string{
{{range $key, $value := .comments}}
{{$key}}: "{{$value}}",{{end}}
}

// GetMsg get error msg
func GetMsg(code int) string {
var (
msg string
ok bool
)
if msg, ok = messages[code]; !ok {
msg = noErrorMsg
}
return msg
}
`

// gen 生成代码
func gen(comments map[string]string) ([]byte, error) {
var buf = bytes.NewBufferString("")

data := map[string]interface{}{
"pkg": os.Getenv("GOPACKAGE"),
"comments": comments,
}

t, err := template.New("").Parse(tpl)
if err != nil {
return nil, errors.Wrapf(err, "template init err")
}

err = t.Execute(buf, data)
if err != nil {
return nil, errors.Wrapf(err, "template data err")
}

return format.Source(buf.Bytes())
}

总结

从一个简单的效率需求引申到go generateast的使用,顺便阅读了一下ast的源码,花费的时间其实可能是这个工具节约的时间的几倍了,但是收获也是之前没有想到的。

  1. 使用了这么久的go命令,详细的阅读了go help command的说明之后,发现之前可能连了解都算不上
  2. 标准库的godoc是最好的使用说明,第二好的是它的源代码

参考资料

  1. go-const-msg 本文实现的源代码

  2. Golang Generate 命令说明与使用

  3. AST 标准库文档

  4. Go 的 AST(抽象语法树)

  5. GO 官方博客: Generating code

  • 本文作者: mohuishou <1@lailin.xyz>
  • 本文链接: https://lailin.xyz/post/41140.html
  • 版权声明: 本博客所有文章除特别声明外,均采用 BY-NC-SA 许可协议。转载请注明出处!