github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/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 = "_local.go"
    20  	buildTag   = "local"
    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  		}
    54  
    55  		var changed bool
    56  		var err error
    57  		for _, fnc := range fncs {
    58  			var c bool
    59  			c, err = fnc(prog.Fset, pi.Info, qual, f)
    60  			if c {
    61  				changed = true
    62  			}
    63  			if err != nil {
    64  				break
    65  			}
    66  		}
    67  
    68  		if err != nil {
    69  			fmt.Fprintf(os.Stderr, "could not rewrite %v: %v\n", f, err)
    70  			os.Exit(2)
    71  		}
    72  
    73  		if !changed {
    74  			continue
    75  		}
    76  
    77  		fmt.Println(prog.Fset.Position(f.Pos()).Filename)
    78  
    79  		origHasBuildTag := false
    80  
    81  		for _, c := range f.Comments {
    82  			for _, c := range c.List {
    83  				if c.Text == "// +build !"+buildTag {
    84  					c.Text = "// +build " + buildTag
    85  					origHasBuildTag = true
    86  				}
    87  			}
    88  		}
    89  
    90  		var buf bytes.Buffer
    91  		fpath := prog.Fset.Position(f.Pos()).Filename
    92  		if origHasBuildTag {
    93  			printer.Fprint(&buf, prog.Fset, f)
    94  		} else {
    95  			buf.Write([]byte("// +build " + buildTag + "\n\n"))
    96  			printer.Fprint(&buf, prog.Fset, f)
    97  
    98  			// prepend build comment to original file
    99  			b, err := ioutil.ReadFile(fpath)
   100  			if err != nil {
   101  				fmt.Fprintf(os.Stderr, "could not read source file: %v\n", err)
   102  				os.Exit(2)
   103  			}
   104  			b = append([]byte("// +build !"+buildTag+"\n\n"), b...)
   105  			b, err = format.Source(b)
   106  			if err != nil {
   107  				fmt.Fprintf(os.Stderr, "could not format source file %v: %v\n", filepath.Base(fpath), err)
   108  				os.Exit(2)
   109  			}
   110  			f, err := os.OpenFile(fpath, os.O_WRONLY, 0)
   111  			if err != nil {
   112  				fmt.Fprintf(os.Stderr, "could not open source file for writing: %v\n", err)
   113  				os.Exit(2)
   114  			}
   115  			if _, err = f.Write(b); err != nil {
   116  				fmt.Fprintf(os.Stderr, "could not write to source file: %v\n", err)
   117  				os.Exit(2)
   118  			}
   119  		}
   120  
   121  		b, err := format.Source(buf.Bytes())
   122  		if err != nil {
   123  			panic(fmt.Errorf("unexpected internal error: %v", err))
   124  		}
   125  		fpath = fpath[:len(fpath)-3] + fileSuffix
   126  		if err = ioutil.WriteFile(fpath, b, 0664); err != nil {
   127  			fmt.Fprintf(os.Stderr, "could not create instrument source file: %v\n", err)
   128  			os.Exit(2)
   129  		}
   130  	}
   131  }