go generate and ast-蒲公英云

原文首发于我的博客: lailin.xyz/post/41140.…

楔(xiē)子

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

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

不写，就满屏波浪线，不能忍！

写了，就得Copy一份，还不利于维护，不能忍！

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

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

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

? go generate

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

? 命令文档

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

▶ go help generate
usage: go generate [-run regexp] [-n] [-v] [-x] [build flags] [file.go... | packages]
... 
复制代码

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

参数说明
- -run 正则表达式匹配命令行，仅执行匹配的命令(和go test -run类似)
- -v 打印已被检索处理的文件。
- -n 打印出将被执行的命令，此时将不真实执行命令
- -x 打印已执行的命令

举个栗子

# 对当前包下的Go文件进行处理, 并打印已被检索处理的文件。
go generate -v 
# 打印当前目录下所有文件中将要被执行的命令(实际不会执行)
go generate -n ./...
复制代码

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

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

# PS: 这是一个正则表达式
^// Code generated .* DO NOT EDIT\.$
复制代码

举个栗子:

// Code generated by mohuishou DO NOT EDIT
package painkiller
复制代码

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

$GOARCH
    系统架构: arm, amd64 等
$GOOS
    操作系统: linux, windows 等
$GOFILE
    当前执行的命令所处的文件名
$GOLINE
    当前执行的命令在文件中的行号
$GOPACKAGE
    执行的命令所处的文件的包名
$DOLLAR
    $ 符号
复制代码

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

源文件: painkiller.go

//go:generate stringer -type=Pill
package painkiller
type Pill int
const (
    Placebo Pill = iota
    Aspirin
    Ibuprofen
    Paracetamol
    Acetaminophen = Paracetamol
)
复制代码

执行命令

go generate
复制代码

生成文件: painkiller_stringer.go

// 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包

基础的接口类型

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

// 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例子

// CommentMap把AST节点和其关联的注释列表进行映射
type CommentMap map[Node][]*CommentGroup
复制代码

通过parse读取源码创建一个AST

fset := token.NewFileSet() // positions are relative to fset
f, err := parser.ParseFile(fset, "src.go", src, parser.ParseComments)
if err != nil {
    panic(err)
}
复制代码

从AST中新建一个CommentMap

cmap := ast.NewCommentMap(fset, f, f.Comments)
复制代码

需求实现

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

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. 获取注释信息

// 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. 生成代码

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