forked from charithe/go-ast-refactoring
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
123 lines (98 loc) · 2.73 KB
/
main.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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package main
import (
"fmt"
"go/ast"
"go/token"
"log"
"os"
"path/filepath"
"github.com/dave/dst"
"github.com/dave/dst/decorator"
"github.com/dave/dst/decorator/resolver/goast"
"github.com/dave/dst/decorator/resolver/guess"
"golang.org/x/tools/go/packages"
)
func main() {
// find the absolute path to the example.
dir, err := filepath.Abs("/home/cell/code/cloud-on-k8s/")
if err != nil {
exitOnErr(err)
}
// load the packages
fset := token.NewFileSet()
cfg := &packages.Config{
Mode: packages.NeedTypes | packages.NeedTypesInfo | packages.NeedFiles | packages.NeedSyntax | packages.NeedName | packages.NeedImports | packages.NeedDeps,
Dir: dir,
Fset: fset,
Logf: log.Printf,
Tests: true,
}
pkgs, err := packages.Load(cfg, "./...")
exitOnErr(err)
// create a new decorator with support for resolving imports
//d := decorator.NewDecoratorWithImports(fset, "main", goast.New())
d := decorator.NewDecoratorWithImports(fset, "main", goast.WithResolver(guess.WithMap(map[string]string{"github.com/elastic/cloud-on-k8s/pkg/utils/log": "ulog"})))
// rewrite sources
for _, pkg := range pkgs {
for _, f := range pkg.Syntax {
rewrite(d, pkg, f, iface)
}
}
}
func rewrite(d *decorator.Decorator, pkg *packages.Package, f *ast.File) {
updated := false
df, err := d.DecorateFile(f)
exitOnErr(err)
// traverse the AST.
dst.Inspect(df, func(node dst.Node) bool {
if node == nil {
return false
}
ve, ok := node.(*dst.ValueSpec)
if !ok || len(ve.Values) != 1 {
return true
}
ce, ok := ve.Values[0].(*dst.CallExpr)
if !ok {
return true
}
se, ok := ce.Fun.(*dst.SelectorExpr)
if !ok {
return true
}
//dst.Fprint(os.Stdout, se, nil)
ident, ok := se.X.(*dst.Ident)
if !ok {
return true
}
if ident.Path == "sigs.k8s.io/controller-runtime/pkg/log" && ident.Name == "Log" {
log.Printf("Match found at %s", d.Fset.Position(d.Ast.Nodes[se].Pos()))
newIdent := dst.NewIdent("Log")
newIdent.Path = "github.com/elastic/cloud-on-k8s/pkg/utils/log"
se.X = newIdent
updated = true
}
return true
})
if updated {
p := d.Fset.Position(f.Pos())
writeFile(p.Filename, df)
}
}
func writeFile(fileName string, df *dst.File) {
log.Printf("Writing changes to %s", fileName)
restorer := decorator.NewRestorerWithImports("main", guess.WithMap(map[string]string{"github.com/elastic/cloud-on-k8s/pkg/utils/log": "ulog"}))
out, err := os.Create(fileName)
if err != nil {
exitOnErr(fmt.Errorf("failed to open %s for writing", fileName, err))
}
defer out.Close()
fr := restorer.FileRestorer()
fr.Alias["github.com/elastic/cloud-on-k8s/pkg/utils/log"] = "ulog"
exitOnErr(fr.Fprint(out, df))
}
func exitOnErr(err error) {
if err != nil {
log.Fatalf("ERROR: %v\n", err)
}
}