48 lines
1017 B
Go
48 lines
1017 B
Go
|
package ast
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"go/ast"
|
||
|
|
"go/token"
|
||
|
|
)
|
||
|
|
|
||
|
|
// 增加 import 方法
|
||
|
|
func AddImport(astNode ast.Node, imp string) {
|
||
|
|
impStr := fmt.Sprintf("\"%s\"", imp)
|
||
|
|
ast.Inspect(astNode, func(node ast.Node) bool {
|
||
|
|
if genDecl, ok := node.(*ast.GenDecl); ok {
|
||
|
|
if genDecl.Tok == token.IMPORT {
|
||
|
|
for i := range genDecl.Specs {
|
||
|
|
if impNode, ok := genDecl.Specs[i].(*ast.ImportSpec); ok {
|
||
|
|
if impNode.Path.Value == impStr {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
genDecl.Specs = append(genDecl.Specs, &ast.ImportSpec{
|
||
|
|
Path: &ast.BasicLit{
|
||
|
|
Kind: token.STRING,
|
||
|
|
Value: impStr,
|
||
|
|
},
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
// 查询特定function方法
|
||
|
|
func FindFunction(astNode ast.Node, FunctionName string) *ast.FuncDecl {
|
||
|
|
var funcDeclP *ast.FuncDecl
|
||
|
|
ast.Inspect(astNode, func(node ast.Node) bool {
|
||
|
|
if funcDecl, ok := node.(*ast.FuncDecl); ok {
|
||
|
|
if funcDecl.Name.String() == FunctionName {
|
||
|
|
funcDeclP = funcDecl
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
})
|
||
|
|
return funcDeclP
|
||
|
|
}
|