利用 go-ast 語法樹做代碼生成
【導讀】本文中作者記錄了用 go/ast 封裝 zap 日誌庫做中間件的操作,思路流暢,學起來!
需求概述
go.uber.org/zap
日誌包性能很好,但是用起來很不方便,雖然新版本添加了 global 方法,但仍然彆扭:zap.S().Info()
。
現在我們的需求就是將 zap 的 sugaredLogger 封裝成一個包,讓它像 logrus
一樣易用,直接調用包內函數:log.Info()
。
我們只需要找到 `SugaredLogger 這個 type 擁有的 Exported 方法,將其改爲函數,函數體調用其同名方法:
func Info(args ...interface{}) {
_globalS.Info(args)
}
此處 var _globalS = zap.S()
,因爲 zap.S()
每次調用都會調用 RWMutex.RLock()
,改爲全局變量提高性能。
這個需求很簡單,黏貼複製一頓 replace 就可以搞定,但這太蠢,我們要用一種更 Geek 的方式:代碼生成。
完整代碼:https://github.com/win5do/go-lib/blob/edc6813f5414f1251e91b670c3a9b89ed89e3525/logx/generator/zap.go
代碼實現
要獲取某個 type 的方法,大家可能會想到 reflect
反射包,但是 reflect 只能知道參數類型,沒法知道參數名。所以這裏我們使用go/ast
直接解析源碼。
獲取 ast 語法樹
方法可能分散在包內不同 go 文件,所以必須解析整個包,而不是單個文件。
首先要找到 go.uber.org/zap
的源碼路徑,這裏我們極客到底,通過 go/build 包獲取其在 gomod 中的路徑,不用手動填寫:
func getImportPkg(pkg string) (string, error) {
p, err := gobuild.Import(pkg, "", gobuild.FindOnly)
if err != nil {
return "", err
}
return p.Dir, err
}
解析整個 zap 包,拿到 ast 語法樹:
func parseDir(dir, pkgName string) (*ast.Package, error) {
pkgMap, err := goparser.ParseDir(
token.NewFileSet(),
dir,
func(info os.FileInfo) bool {
// skip go-test
return !strings.Contains(info.Name(), "_test.go")
},
goparser.Mode(0), // no comment
)
if err != nil {
return nil, errx.WithStackOnce(err)
}
pkg, ok := pkgMap[pkgName]
if !ok {
err := errors.New("not found")
return nil, errx.WithStackOnce(err)
}
return pkg, nil
}
遍歷並修改 ast
遍歷 ast,找到 SugaredLogger 的所有 Exported 方法:
func (v *visitor) Visit(node ast.Node) ast.Visitor {
switch n := node.(type) {
case *ast.FuncDecl:
if n.Recv == nil ||
!n.Name.IsExported() ||
len(n.Recv.List) != 1 {
return nil
}
t, ok := n.Recv.List[0].Type.(*ast.StarExpr)
if !ok {
return nil
}
if t.X.(*ast.Ident).String() != "SugaredLogger" {
return nil
}
log.Printf("func name: %s", n.Name.String())
v.funcs = append(v.funcs, rewriteFunc(n))
}
return v
}
func rewriteFunc(fn *ast.FuncDecl) *ast.FuncDecl {
fn.Recv = nil
fnName := fn.Name.String()
var args []string
for _, field := range fn.Type.Params.List {
for _, id := range field.Names {
idStr := id.String()
_, ok := field.Type.(*ast.Ellipsis)
if ok {
// Ellipsis args
idStr += "..."
}
args = append(args, idStr)
}
}
exprStr := fmt.Sprintf(`_globalS.%s(%s)`, fnName, strings.Join(args, ","))
expr, err := goparser.ParseExpr(exprStr)
if err != nil {
panic(err)
}
var body []ast.Stmt
if fn.Type.Results != nil {
body = []ast.Stmt{
&ast.ReturnStmt{
// Return:
Results: []ast.Expr{expr},
},
}
} else {
body = []ast.Stmt{
&ast.ExprStmt{
X: expr,
},
}
}
fn.Body.List = body
return fn
}
上一步函數返回值中 zap.SugaredLogger
在目標包中需要改爲 zap.SugaredLogger
,這裏使用 type alias 簡單處理一下,當然修改 ast 同樣能做到:
// alias
type (
Logger = zap.Logger
SugaredLogger = zap.SugaredLogger
)
ast 轉化爲 go 代碼
單個 func 的 ast 轉化爲 go 代碼,使用 go/format
包:
func astToGo(dst *bytes.Buffer, node interface{}) error {
addNewline := func() {
err := dst.WriteByte('\n') // add newline
if err != nil {
log.Panicln(err)
}
}
addNewline()
err := format.Node(dst, token.NewFileSet(), node)
if err != nil {
return err
}
addNewline()
return nil
}
拼裝成完整 go file:
func writeGoFile(wr io.Writer, funcs []ast.Decl) error {
// 輸出Go代碼
header := `// Code generated by log-gen. DO NOT EDIT.
package log
`
buffer := bytes.NewBufferString(header)
for _, fn := range funcs {
err := astToGo(buffer, fn)
if err != nil {
return errx.WithStackOnce(err)
}
}
_, err := wr.Write(buffer.Bytes())
return err
}
這個程序是輸出到了 os.Stdout,通過 go:generate
將其重定向到 zap_sugar_generated.go 文件中:
//go:generate sh -c "go run ./generator >zap_sugar_generated.go"
大功告成,輸出代碼示例:
// Code generated by log-gen. DO NOT EDIT.
package log
func Desugar() *Logger {
return _globalS.Desugar()
}
func Named(name string) *SugaredLogger {
return _globalS.Named(name)
}
func With(args ...interface{}) *SugaredLogger {
return _globalS.With(args...)
}
func Debug(args ...interface{}) {
_globalS.Debug(args...)
}
func Info(args ...interface{}) {
_globalS.Info(args...)
}
func Warn(args ...interface{}) {
_globalS.Warn(args...)
}
func Error(args ...interface{}) {
_globalS.Error(args...)
}
// ......
即使之後 zap 包升級了,方法有增改,修改 gomod 版本再次執行 gernerate 即可一鍵同步,告別手動復粘。
總結
Go 沒法像 Java 那樣做動態 AOP,但可以通過 go/ast 做代碼生成,達成同樣目標,而且不像 reflect 會影響性能和靜態檢查。用的好的話可以極大提高效率,更加自動化,減少手工復粘,也就降低犯錯概率。
已在很多明星開源項目裏廣泛應用,如:
-
代碼編輯工具 gomodifytags:https://github.com/fatih/gomo...、
-
Go 編譯時依賴注入 Wire:https://github.com/google/wire
-
K8S 源碼:https://github.com/kubernetes...
轉自:
segmentfault.com/a/1190000039215176
Go 開發大全
參與維護一個非常全面的 Go 開源技術資源庫。日常分享 Go, 雲原生、k8s、Docker 和微服務方面的技術文章和行業動態。
本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源:https://mp.weixin.qq.com/s/qqzmcn6D7IjSUAT3VO6Sbg