-
Notifications
You must be signed in to change notification settings - Fork 3
/
util.go
99 lines (85 loc) · 2.39 KB
/
util.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
package logannotation
import (
"fmt"
"go/ast"
"go/parser"
)
// NewIdents 按顺序生成每个名称的 ast.Ident
func NewIdents(names ...string) []*ast.Ident {
idents := make([]*ast.Ident, len(names))
for i, name := range names {
idents[i] = ast.NewIdent(name)
}
return idents
}
// GetIdentNames 按顺序返回 idents 的名称列表
func GetIdentNames(idents ...*ast.Ident) []string {
names := make([]string, len(idents))
for i, ident := range idents {
names[i] = ident.Name
}
return names
}
// 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
}
// GetFieldNames 按顺序返回 fields 的名称列表
// 如果一个 Field 有多个名称,那么它们会在一起
// (a int, b, c string, d bool) 会返回 ("a", "b", "c", "d")
func GetFieldNames(fields ...*ast.Field) []string {
var names []string
for _, field := range fields {
names = append(names, GetIdentNames(field.Names...)...)
}
return names
}
// 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
}
// NewFuncLitDefer 产生一个 defer 语句,运行一个匿名函数,函数体是入参语句列表
func NewFuncLitDefer(funcStmts ...ast.Stmt) *ast.DeferStmt {
return &ast.DeferStmt{
Call: &ast.CallExpr{
Fun: NewFuncLit(&ast.FuncType{}, funcStmts...),
},
}
}
// NewFuncLit 产生一个匿名函数
func NewFuncLit(funcType *ast.FuncType, stmts ...ast.Stmt) *ast.FuncLit {
return &ast.FuncLit{
Type: funcType,
Body: &ast.BlockStmt{
List: stmts,
},
}
}