利用 go/ast 語法樹作代碼生成

時間 2021-02-17

標籤 node git github app 函數工具 post 性能 google 欄目 Git 简体版

原文原文鏈接

需求概述

go.uber.org/zap 日誌包性能很好，可是用起來很不方便，雖然新版本添加了 global 方法，但仍然彆扭：zap.S().Info()。node

如今咱們的需求就是將 zap 的 sugaredLogger 封裝成一個包，讓它像 logrus 同樣易用，直接調用包內函數：log.Info()。git

咱們只須要找到`SugaredLogger這個 type 擁有的 Exported 方法，將其改成函數，函數體調用其同名方法：github

func Info(args ...interface{}) {
    _globalS.Info(args)
}

此處 var _globalS = zap.S()，由於 zap.S() 每次調用都會調用 RWMutex.RLock() ，改成全局變量提升性能。app

這個需求很簡單，黏貼複製一頓 replace 就能夠搞定，但這太蠢，咱們要用一種更 Geek 的方式：代碼生成。函數

完整代碼：https://github.com/win5do/go-...工具

代碼實現

要獲取某個 type 的方法，你們可能會想到 reflect 反射包，可是 reflect 只能知道參數類型，無法知道參數名。因此這裏咱們使用go/ast直接解析源碼。post

獲取 ast 語法樹

方法可能分散在包內不一樣 go 文件，因此必須解析整個包，而不是單個文件。性能

首先要找到 go.uber.org/zap 的源碼路徑，這裏咱們極客到底，經過 go/build 包獲取其在 gomod 中的路徑，不用手動填寫：ui

func getImportPkg(pkg string) (string, error) {
    p, err := gobuild.Import(pkg, "", gobuild.FindOnly)
    if err != nil {
        return "", err
    }

    return p.Dir, err
}

解析整個 zap 包，拿到 ast 語法樹：google

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
}

將方法 Recv 置空，變爲函數，參數不變，函數 body 改成調用全局變量 _globalS 的同名方法，若是有返回值則須要 return 語句：

func rewriteFunc(fn *ast.FuncDecl) *ast.FuncDecl {
    fn.Recv = nil

    fnName := fn.Name.String()

    var args []string
    for _, v := range fn.Type.Params.List {
        for _, w := range v.Names {
            args = append(args, w.String())
        }
    }

    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)
}

func DPanic(args ...interface{}) {
    _globalS.DPanic(args)
}

func Panic(args ...interface{}) {
    _globalS.Panic(args)
}

func Fatal(args ...interface{}) {
    _globalS.Fatal(args)
}

// ......

即便以後 zap 包升級了，方法有增改，修改 gomod 版本再次執行 gernerate 便可一鍵同步，告別手動復粘。