✨ Init

utils/ast/ast.go

@@ -0,0 +1,47 @@
+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
+}