-
-
Notifications
You must be signed in to change notification settings - Fork 6k
/
ast_auto_enter.go
47 lines (45 loc) · 990 Bytes
/
ast_auto_enter.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
package ast
import (
"bytes"
"fmt"
"go/ast"
"go/parser"
"go/printer"
"go/token"
"os"
)
func ImportForAutoEnter(path string, funcName string, code string) {
src, err := os.ReadFile(path)
if err != nil {
fmt.Println(err)
}
fileSet := token.NewFileSet()
astFile, err := parser.ParseFile(fileSet, "", src, 0)
ast.Inspect(astFile, func(node ast.Node) bool {
if typeSpec, ok := node.(*ast.TypeSpec); ok {
if typeSpec.Name.Name == funcName {
if st, ok := typeSpec.Type.(*ast.StructType); ok {
for i := range st.Fields.List {
if t, ok := st.Fields.List[i].Type.(*ast.Ident); ok {
if t.Name == code {
return false
}
}
}
sn := &ast.Field{
Type: &ast.Ident{Name: code},
}
st.Fields.List = append(st.Fields.List, sn)
}
}
}
return true
})
var out []byte
bf := bytes.NewBuffer(out)
err = printer.Fprint(bf, fileSet, astFile)
if err != nil {
return
}
_ = os.WriteFile(path, bf.Bytes(), 0666)
}