Golang 中使用“注释注解”生成日志打印逻辑

背景

errerr

// @Log()@Log

编写代码

go/ast

[1]

.go

filepath.WalkwalkFnwalkFn.go

// walkFn 函数会对每个 .go 文件处理，并调用注解处理器func walkFn(path string, info os.FileInfo, err error) error {    // 如果是文件夹，或者不是 .go 文件，则直接返回不处理    if info.IsDir() || !strings.HasSuffix(info.Name(), ".go") {        return nil    }
    // 将 .go 文件解析成 AST    fileSet, file, err := parseFile(path)    // 如果注解修改了内容，则需要生成新的代码    if logannotation.Overwrite(path, fileSet, file) {        buf := &bytes.Buffer{}        if err := format.Node(buf, fileSet, file); err != nil {            panic(err)        }
        // 如果不需要替换，则生成到另一个文件        if !replace {            lastSlashIndex := strings.LastIndex(path, "/")            genDirPath :=path[:lastSlashIndex] + "/_gen/"            if err := os.Mkdir(genDirPath, 0755); err != nil && os.IsNotExist(err){                panic(err)            }            path = genDirPath + path[lastSlashIndex+1:]        }
        if err := ioutil.WriteFile(path, buf.Bytes(), info.Mode()); err != nil {            panic(err)        }    }
    return nil}

遍历 AST

astutil.Applyimport

// Overwrite 会对每个 file 处理，运行注册的注解 handler ，并返回其是否被修改func Overwrite(filepath string, fileSet *token.FileSet, file *ast.File) (modified bool) {    // 初始化处理本次文件所需的信息对象    info := &Info{        Filepath: filepath,        NamedImportAdder: func(name string, path string) bool {            return astutil.AddNamedImport(fileSet, file, name, path)        },    }
    // 遍历当前文件 ast 上的所有节点    astutil.Apply(file, nil, func(cursor *astutil.Cursor) bool {        // 处理 log 注解        info.Node = cursor.Node()        nodeModified, err := Handler.Handle(info)        if err != nil {            panic(err)        }
        if nodeModified {            modified = nodeModified        }
        return true    })
    return}

识别注释注解

@Log()

func (h *handler) Handle(info *Info) (modified bool, err error) {    // log 注解只用于函数    funcDecl, ok := info.Node.(*ast.FuncDecl)    if !ok {        return    }
    // 如果没有注释，则直接处理下一个    if funcDecl.Doc == nil {        return    }
    // 如果不是可以处理的注解，则直接返回    doc := strings.Trim(funcDecl.Doc.Text(), "\t \n")    if doc != "@Log()" {        return    }    ...}

获取函数入参和出参

funcDecl.Type.Results_0_1...

// SetDefaultNames 给没有名字的 Field 设置默认的名称// 默认名称格式：_0, _1, ...// true: 表示至少设置了一个名称// false: 表示未设置过名称func SetDefaultNames(fields ...*ast.Field) bool {    index := 0    for _, field := range fields {        if field.Names == nil {            field.Names = NewIdents(fmt.Sprintf("_%v", index))            index++        }    }
    return index > 0}

获取打印语句

log.Logger.WithContext(ctx).WithField("filepath", filepath).Infof(format, arg0, arg1)parser.ParseExpr(format, arg0, arg1)

// NewCallExpr 产生一个调用表达式// 待产生表达式：log.Logger.WithContext(ctx).Infof(arg0, arg1)// 其中：//    funcSelector = "log.Logger.WithContext(ctx).Infof"//     args = ("arg0", "arg1")// 调用语句：NewCallExpr("log.Logger.WithContext(ctx).Infof", "arg0", "arg1")func NewCallExpr(funcSelector string, args ...string) (*ast.CallExpr, error) {    // 获取函数对应的表达式    funcExpr, err := parser.ParseExpr(funcSelector)    if err != nil {        return nil, err    }
    // 组装参数列表    argsExpr := make([]ast.Expr, len(args))    for i, arg := range args {        argsExpr[i] = ast.NewIdent(arg)    }
    return &ast.CallExpr{        Fun:  funcExpr,        Args: argsExpr,    }, nil}

defer

// NewFuncLitDefer 产生一个 defer 语句，运行一个匿名函数，函数体是入参语句列表func NewFuncLitDefer(funcStmts ...ast.Stmt) *ast.DeferStmt {    return &ast.DeferStmt{        Call: &ast.CallExpr{            Fun: NewFuncLit(&ast.FuncType{}, funcStmts...),        },    }}

修改函数体

至此我们已经获得了打印入参和出参的语句，接下来就是把他们放在原本函数体的最前面，保证开始和结束时执行。

toBeAddedStmts := []ast.Stmt{    &ast.ExprStmt{X: beforeExpr},    // 离开函数时的语句使用 defer 调用    NewFuncLitDefer(&ast.ExprStmt{X: afterExpr}),}
// 我们将添加的语句放在函数体最前面funcDecl.Body.List = append(toBeAddedStmts, funcDecl.Body.List...)

运行

为了测试我们的注释注解是否工作正确，我们使用如下代码进行测试：

package main
import (    "context"    "logannotation/testdata/log")
func main() {    fn(context.Background(), 1, "2", "3", true)}
// @Log()func fn(ctx context.Context, a int, b, c string, d bool) (int, string, string) {    log.Logger.WithContext(ctx).Infof("#fn executing...")    return a, b, c}

go run logannotation/cmd/generator /Users/idealism/Workspaces/Go/golang-log-annotation/testdata/Users/idealism/Workspaces/Go/golang-log-annotation/testdata/_gen

package main
import (    "context"    "logannotation/testdata/log")
func main() {    fn(context.Background(), 1, "2", "3", true)}
func fn(ctx context.Context, a int, b, c string, d bool) (_0 int, _1 string, _2 string) {    log.Logger.WithContext(        ctx).WithField("filepath", "/Users/idealism/Workspaces/Go/golang-log-annotation/testdata/main.go").Infof("#fn start, params: %+v, %+v, %+v, %+v", a, b, c, d)    defer func() {        log.Logger.WithContext(            ctx).WithField("filepath", "/Users/idealism/Workspaces/Go/golang-log-annotation/testdata/main.go").Infof("#fn end, results: %+v, %+v, %+v", _0, _1, _2)    }()
    log.Logger.WithContext(ctx).Infof("#fn executing...")    return a, b, c}

ctxtraceIdlogrusHooksFields

INFO[0000]#fnstart,params:1,2,3,truefilepath=/Users/idealism/Workspaces/Go/golang-log-annotation/testdata/main.goINFO[0000] #fn executing...                             INFO[0000] #fn end, results: 1, 2, 3                     filepath=/Users/idealism/Workspaces/Go/golang-log-annotation/testdata/main.go

扩展

以上代码是一种简单方式地定制化处理注解，仅处理了打印日志这一逻辑，当然还存在更多扩展的可能性和优化。

•注册自定义注解（这样可以把更多重复逻辑抽出来，例如：参数校验、缓存等逻辑）•同时使用多个注解•注解解析成语法树，支持注解参数•生成的代码仅在需要时换行

Demo

[2]

References

[1][2]