go语言dst库简单应用
dst是一个基于go语言实现的，操纵ast的第三方库。
文件解析、修改与写入的一般流程：
content, _ := os.ReadFile(root_path)
f, _ := decorator.ParseFile(nil, root_path, content, parser.ParseComments) // 解析文件
dstutil.Apply(f, func(cursor *dstutil.Cursor) bool { // 操作ast树
FilterAndEdit(f, cursor) // 自定义实现
return true
}, func(cursor *dstutil.Cursor) bool {
return true
})
var buf bytes.Buffer
writer := io.Writer(&buf)
fset, af, _ := decorator.RestoreFile(f)
printer.Fprint(writer, fset, af) // 将文件格式化输出到writer中
fileContent := buf.String()
os.WriteFile("./other/temp.go", []byte(fileContent), 0666) // 将修改后的内容写入文件
这段代码中dstutil.Apply方法是关键，这个函数可以递归遍历语法树，接收根节点，pre函数与post函数作为参数。pre函数的调用时机是第一次到达节点还没有遍历其子孙节点，如果函数返回false，停止遍历子结点且post函数也不会调用，post函数则是已经遍历过子孙节点回溯到该节点时调用，如果返回false，停止递归。pre与post可以类比前序遍历与后序遍历。
解析字符串
除了文件外，dst提供了直接解析字符串的api
func GoStringToStats(goString string) []dst.Stmt {
parsed, err := decorator.Parse(fmt.Sprintf(`
package main
func main() {
%s
}`, goString))
if err != nil {
panic(fmt.Sprintf("parsing go failure: %v\n%s", err, goString))
}
return parsed.Decls[0].(*dst.FuncDecl).Body.List
}
上面这个例子提供了一个dst库的使用小技巧，可以将任意合法字符串转化为Stmt。
解析与修改声明：
在go语言文件中我们一定会做声明操作，dst库将go语言中的各种声明划分为三类：
type (
// A BadDecl node is a placeholder for a declaration containing
// syntax errors for which a correct declaration node cannot be
// created.
//
BadDecl struct {
Length int // position range of bad declaration
Decs BadDeclDecorations
}
// A GenDecl node (generic declaration node) represents an import,
// constant, type or variable declaration. A valid Lparen position
// (Lparen.IsValid()) indicates a parenthesized declaration.
//
// Relationship between Tok value and Specs element type:
//
// token.IMPORT *ImportSpec
// token.CONST *ValueSpec
// token.TYPE *TypeSpec
// token.VAR *ValueSpec
//
GenDecl struct {
Tok token.Token // IMPORT, CONST, TYPE, or VAR
Lparen bool
Specs []Spec
Rparen bool
Decs GenDeclDecorations
}
// A FuncDecl node represents a function declaration.
FuncDecl struct {
Recv *FieldList // receiver (methods); or nil (functions)
Name *Ident // function/method name
Type *FuncType // function signature: type and value parameters, results, and position of "func" keyword
Body *BlockStmt // function body; or nil for external (non-Go) function
Decs FuncDeclDecorations
}
)
这三个类都实现了了Decl接口, 最常用的是GenDecl与FuncDecl，其中GenDecl可以表示对常量与变量的声明、对import的声明、对结构体和接口的声明。FuncDecl表示对函数的声明。
type Decl interface {
Node
declNode()
}
例1：接下来这个例子演示了如何为文件添加一个int类型的变量声明：
func CustomizedEnhance(file *dst.File, cursor *dstutil.Cursor) {
file.Decls = append(file.Decls, &dst.GenDecl{ // Decls变量存储了文件中的所有顶层声明
Tok: token.VAR, // token.Var 表示声明一个变量, 可以替换为token.Type等来尝试其它类型的声明
Specs: []dst.Spec{
&dst.ValueSpec{
Names: []*dst.Ident{dst.NewIdent("_xgo_res_2")},
Type: dst.NewIdent("int"),
},
},
})
}
dst.File的Decls保存了一个go语言文件中所有的声明并以切片的形式保存，我们可以向切片中添加元素来添加声明。
例2：操作FuncDecl
首先来回顾一下FuncDecl的结构
FuncDecl struct {
Recv *FieldList // receiver (methods); or nil (functions)
Name *Ident // function/method name
Type *FuncType // function signature: type and value parameters, results, and position of "func" keyword
Body *BlockStmt // function body; or nil for external (non-Go) function
Decs FuncDeclDecorations
}
关键属性Type，其类型为FuncType，保存了函数入参与返回值的信息。Body属性则保存了函数的函数体。
FuncType struct {
Func bool
TypeParams *FieldList // type parameters; or nil
Params *FieldList // (incoming) parameters; non-nil
Results *FieldList // (outgoing) results; or nil
Decs FuncTypeDecorations
}
Params保存了函数的所有入参编译期信息，同理Result保存了返回值的编译期信息，我们可以通过修改这两个属性修改入参与出参的数量、名称、类型等基本信息。
为函数添加defer语句
func EnhanceFunc(fun *dst.FuncDecl) {
deferStmt := &dst.DeferStmt{
Call: &dst.CallExpr{
Fun: dst.NewIdent("println"),
Args: []dst.Expr{dst.NewIdent("\"defer\"")},
},
}
fun.Body.List = append([]dst.Stmt{deferStmt}, fun.Body.List...)
}
// 原函数
func (buf *Buf) checkModify(arg1 int, arg2 string) string{
return ""
}
// 修改后函数
func (buf *Buf) checkModify(arg1 int, arg2 string) string{
defer println("defer")
return ""
}
为函数添加更复杂的defer函数
func EnhanceFunc(fun *dst.FuncDecl, code string) {
fun.Body.Decs.Lbrace.Prepend(code)
body := &dst.BlockStmt{
List: []dst.Stmt{
&dst.ExprStmt{
X: &dst.CallExpr{
Fun: dst.NewIdent("fmt.Println"),
Args: []dst.Expr{dst.NewIdent("\"defer1\"")},
},
},
&dst.ExprStmt{
X: &dst.CallExpr{
Fun: dst.NewIdent("println"),
Args: []dst.Expr{dst.NewIdent("\"defer2\"")},
},
},
},
}
deferStmt := &dst.DeferStmt{
Call: &dst.CallExpr{
Fun: &dst.FuncLit{
Type: &dst.FuncType{
Params: &dst.FieldList{},
},
Body: body,
},
},
}
fun.Body.List = append([]dst.Stmt{deferStmt}, fun.Body.List...)
}
BlockStat可以容纳多个以及多种类型的Stat，可以将多个语句组合为Body，再进一步将Body组合为Func。FuncLit可以表示一个完整的函数信息，而FuncType只能表示函数签名信息，不能表示函数整体的声明。
dst库还提供了其他种类丰富的Stat，包括但不限于GoStat，对应协程、IfStat，对应if语句块、TypeSwitchStmt对应类型switch语句块等等。
Spec结点：
概念上与Decl似乎有相似之处，但是在实践上还是能发现不同之处。dst库也将它们划分为三类。
type (
// The Spec type stands for any of *ImportSpec, *ValueSpec, and *TypeSpec.
Spec interface {
Node
specNode()
}
// An ImportSpec node represents a single package import.
ImportSpec struct {
Name *Ident // local package name (including "."); or nil
Path *BasicLit // import path
Decs ImportSpecDecorations
}
// A ValueSpec node represents a constant or variable declaration
// (ConstSpec or VarSpec production).
//
ValueSpec struct {
Names []*Ident // value names (len(Names) > 0)
Type Expr // value type; or nil
Values []Expr // initial values; or nil
Decs ValueSpecDecorations
}
// A TypeSpec node represents a type declaration (TypeSpec production).
TypeSpec struct {
Name *Ident // type name
TypeParams *FieldList // type parameters; or nil
Assign bool // position of '=', if any
Type Expr // *Ident, *ParenExpr, *SelectorExpr, *StarExpr, or any of the *XxxTypes
Decs TypeSpecDecorations
}
)
注意ValueSpec中的Names与Values都是数组，这是因为go语言允许类似var a, b int的语法，可以一次性声明多个变量。
封装逻辑
基于前文的铺垫我们可以将修改函数内容与模板字符串联系起来达到更加灵活的效果。
func ExecuteTemplate(tmpl string, data interface{}) string {
t, err := template.New("").Parse(tmpl)
if err != nil {
panic(err)
}
var buf bytes.Buffer
err = t.Execute(&buf, data)
if err != nil {
panic(err)
}
return buf.String()
}
获取模板字符串。 template库是go语言内置的模板字符串生成工具，至于go语言中模板字符串的生成规则可以自行百度搜索。
func InsertStmtsBeforeBody(body *dst.BlockStmt, tmpl string, data interface{}) {
body.List = append(GoStringToStats(ExecuteTemplate(tmpl, data)), body.List...)
}
引入模板字符串我们可以应对更加复杂的字符串拼接与生成任务，更重要的是可以在运行时动态获取字符串。GoStringToStats函数实现在第二个加粗标题处有实现。
参考资料：
dst库地址：https://pkg.go.dev/github.com/dave/dst
浙公网安备 33010602011771号