go generate and ast

時間 2019-11-22

標籤 generate ast 简体版

原文原文鏈接

原文首發於個人博客: lailin.xyz/post/41140.…html

楔(xiē)子

最近寫API CURD比較多，爲告終構清晰，返回值須要統一錯誤碼，因此在一個統一的errcode包中定義錯誤碼常量，以及其錯誤信息.node

以下圖所示，因爲常量是導出字符 -> golint 檢測須要編寫註釋 -> 註釋信息其實就是錯誤信息，已經在下文的msg map[int]string中定義，若是在寫就得寫兩遍linux

不寫，就滿屏波浪線，不能忍！git

寫了，就得Copy一份，還不利於維護，不能忍！github

能不能只寫一份註釋，剩下的msg經過讀取註釋信息自動生成，將咱們寶(hua)貴(diao)的生命，從這些重複繁雜無心義的勞動中解放出來。golang

爲了實現這個~~偉大的目標~~, 須要如下兩個關鍵的數據:正則表達式

解析源代碼獲取常量與註釋之間的關係 -> 🌲Go抽象語法樹: AST^[3]
從Go源碼生成Go代碼 -> 👏 go generate^[5]

👏 go generate

golang在1.4版本中引入了go generate命令，經常使用於文件生成，例如在Golang官方博客^[5]中介紹的Stringer能夠爲枚舉自動實現Stringer的方法，從業務代碼中解放出來windows

💻 命令文檔

使用go help generate咱們能夠查看一下命令的幫助文檔bash

▶ 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())
}
複製代碼