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  }