go generate and ast

你的名字 2022-01-19 23:23 331阅读 0赞

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

## 楔(xiē)子 ##

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

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

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

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

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

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

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

## ? go generate ##

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

### ? 命令文档 ###

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

▶ 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.  举个栗子
    
        # 对当前包下的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`生成的，请在生成的代码中添加注释(一般放在文件开头)
    
        # PS: 这是一个正则表达式
        ^// Code generated .* DO NOT EDIT\.$
        复制代码
    
    举个栗子:
    
        // Code generated by mohuishou DO NOT EDIT
        
        package painkiller
        复制代码
6.  `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\]][3]，Go的AST(抽象语法树)[\[4\]][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\]][3]的Example中可以发现有一个[CommentMap][]例子

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

1.  通过`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)
        }
        复制代码
2.  从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())
    }
    复制代码

## 总结 ##

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

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

## 参考资料 ##

1.  [go-const-msg 本文实现的源代码][go-const-msg]
2.  [Golang Generate命令说明与使用][Golang Generate]
3.  [AST标准库文档][AST]
4.  [Go的AST(抽象语法树)][4]
5.  [GO 官方博客: Generating code][5]

## License ##

*  本文作者： mohuishou [1@lailin.xyz][1_lailin.xyz]
 *  本文链接： [lailin.xyz/post/41140.…][lailin.xyz_post_41140.]
 *  版权声明： 本博客所有文章除特别声明外，均采用 BY-NC-SA 许可协议。转载请注明出处！

[lailin.xyz_post_41140.]: https://link.juejin.im?target=https%3A%2F%2Flailin.xyz%2Fpost%2F41140.html
[3]: https://link.juejin.im?target=https%3A%2F%2Fgolang.org%2Fpkg%2Fgo%2Fast
[5]: https://link.juejin.im?target=https%3A%2F%2Fblog.golang.org%2Fgenerate
[Stringer]: https://link.juejin.im?target=golang.org%2Fx%2Ftools%2Fcmd%2Fstringer
[4]: https://link.juejin.im?target=https%3A%2F%2Fzhuanlan.zhihu.com%2Fp%2F28516587
[CommentMap]: https://link.juejin.im?target=https%3A%2F%2Fgolang.org%2Fpkg%2Fgo%2Fast%2F%23example_CommentMap
[go-const-msg]: https://link.juejin.im?target=https%3A%2F%2Fgithub.com%2Fmohuishou%2Fgo-const-msg
[Golang Generate]: https://link.juejin.im?target=https%3A%2F%2Fyushuangqi.com%2Fblog%2F2017%2Fgo-command-generate.html
[AST]: https://link.juejin.im?target=https%3A%2F%2Fgolang.org%2Fpkg%2Fgo%2Fast%2F
[1_lailin.xyz]: https://link.juejin.im?target=mailto%3A1%40lailin.xyz