github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/other/cmd/rewrite/main.go (about) 1 package main 2 3 import ( 4 "bytes" 5 "fmt" 6 "go/ast" 7 "go/format" 8 "go/printer" 9 "go/token" 10 "go/types" 11 "io/ioutil" 12 "os" 13 "path/filepath" 14 15 "golang.org/x/tools/go/loader" 16 ) 17 18 const ( 19 fileSuffix = "_context.go" 20 buildTag = "context" 21 ) 22 23 func main() { 24 path := "." 25 switch { 26 case len(os.Args) == 2: 27 path = os.Args[1] 28 case len(os.Args) > 2: 29 fmt.Fprintf(os.Stderr, "Usage: %v [<path>]\n", os.Args[0]) 30 os.Exit(1) 31 } 32 33 packageDir(path) 34 } 35 36 func packageDir(dir string) { 37 var conf loader.Config 38 conf.Import(dir) 39 40 prog, err := conf.Load() 41 if err != nil { 42 fmt.Fprintln(os.Stderr, "could not parse source files:", err) 43 os.Exit(2) 44 } 45 46 pi := prog.InitialPackages()[0] 47 48 for _, f := range pi.Files { 49 qual := qualifierForFile(pi.Pkg, f) 50 51 fncs := []func(*token.FileSet, types.Info, types.Qualifier, *ast.File) (bool, error){ 52 rewriteGos, 53 rewriteCalls, 54 } 55 56 var changed bool 57 var err error 58 for _, fnc := range fncs { 59 var c bool 60 c, err = fnc(prog.Fset, pi.Info, qual, f) 61 if c { 62 changed = true 63 } 64 if err != nil { 65 break 66 } 67 } 68 69 if err != nil { 70 fmt.Fprintf(os.Stderr, "could not rewrite %v: %v\n", f, err) 71 os.Exit(2) 72 } 73 74 if !changed { 75 continue 76 } 77 78 fmt.Println(prog.Fset.Position(f.Pos()).Filename) 79 80 origHasBuildTag := false 81 82 for _, c := range f.Comments { 83 for _, c := range c.List { 84 if c.Text == "// +build !"+buildTag { 85 c.Text = "// +build " + buildTag 86 origHasBuildTag = true 87 } 88 } 89 } 90 91 var buf bytes.Buffer 92 fpath := prog.Fset.Position(f.Pos()).Filename 93 if origHasBuildTag { 94 printer.Fprint(&buf, prog.Fset, f) 95 } else { 96 buf.Write([]byte("// +build " + buildTag + "\n\n")) 97 printer.Fprint(&buf, prog.Fset, f) 98 99 // prepend build comment to original file 100 b, err := ioutil.ReadFile(fpath) 101 if err != nil { 102 fmt.Fprintf(os.Stderr, "could not read source file: %v\n", err) 103 os.Exit(2) 104 } 105 b = append([]byte("// +build !"+buildTag+"\n\n"), b...) 106 b, err = format.Source(b) 107 if err != nil { 108 fmt.Fprintf(os.Stderr, "could not format source file %v: %v\n", filepath.Base(fpath), err) 109 os.Exit(2) 110 } 111 f, err := os.OpenFile(fpath, os.O_WRONLY, 0) 112 if err != nil { 113 fmt.Fprintf(os.Stderr, "could not open source file for writing: %v\n", err) 114 os.Exit(2) 115 } 116 if _, err = f.Write(b); err != nil { 117 fmt.Fprintf(os.Stderr, "could not write to source file: %v\n", err) 118 os.Exit(2) 119 } 120 } 121 122 b, err := format.Source(buf.Bytes()) 123 if err != nil { 124 panic(fmt.Errorf("unexpected internal error: %v", err)) 125 } 126 fpath = fpath[:len(fpath)-3] + fileSuffix 127 if err = ioutil.WriteFile(fpath, b, 0664); err != nil { 128 fmt.Fprintf(os.Stderr, "could not create instrument source file: %v\n", err) 129 os.Exit(2) 130 } 131 } 132 }