github.com/chrislusf/greenpack@v3.7.1-0.20170911073826-ad5bd10b7c47+incompatible/printer/print.go (about)

     1  package printer
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"io/ioutil"
     8  	"strings"
     9  
    10  	"github.com/glycerine/greenpack/cfg"
    11  	"github.com/glycerine/greenpack/gen"
    12  	"github.com/glycerine/greenpack/parse"
    13  	"golang.org/x/tools/imports"
    14  )
    15  
    16  func infof(s string, v ...interface{}) {
    17  	fmt.Printf(s, v...)
    18  }
    19  
    20  // PrintFile prints the methods for the provided list
    21  // of elements to the given file name and canonical
    22  // package path.
    23  func PrintFile(
    24  	file string,
    25  	f *parse.FileSet,
    26  	mode gen.Method,
    27  	cfg *cfg.GreenConfig,
    28  	pathToGoSource string) error {
    29  
    30  	out, tests, err := generate(f, mode, cfg)
    31  	if err != nil {
    32  		return err
    33  	}
    34  
    35  	// we'll run goimports on the main file
    36  	// in another goroutine, and run it here
    37  	// for the test file. empirically, this
    38  	// takes about the same amount of time as
    39  	// doing them in serial when GOMAXPROCS=1,
    40  	// and faster otherwise.
    41  	res := goformat(file, out.Bytes())
    42  	if tests != nil {
    43  		testfile := strings.TrimSuffix(file, ".go") + "_test.go"
    44  		err = format(testfile, tests.Bytes())
    45  		if err != nil {
    46  			return err
    47  		}
    48  		infof(">>> Wrote and formatted \"%s\"\n", testfile)
    49  	}
    50  	err = <-res
    51  	if err != nil {
    52  		return err
    53  	}
    54  	return nil
    55  }
    56  
    57  func format(file string, data []byte) error {
    58  	out, err := imports.Process(file, data, nil)
    59  	if err != nil {
    60  		fmt.Printf("\n\n debug: problem file:\n%s\n", file)
    61  		ioutil.WriteFile(file, data, 0600)
    62  		return err
    63  	}
    64  	return ioutil.WriteFile(file, out, 0600)
    65  }
    66  
    67  func goformat(file string, data []byte) <-chan error {
    68  	out := make(chan error, 1)
    69  	go func(file string, data []byte, end chan error) {
    70  		end <- format(file, data)
    71  		infof(">>> Wrote and formatted \"%s\"\n", file)
    72  	}(file, data, out)
    73  	return out
    74  }
    75  
    76  func dedupImports(imp []string) []string {
    77  	m := make(map[string]struct{})
    78  	for i := range imp {
    79  		m[imp[i]] = struct{}{}
    80  	}
    81  	r := []string{}
    82  	for k := range m {
    83  		r = append(r, k)
    84  	}
    85  	return r
    86  }
    87  
    88  func generate(f *parse.FileSet, mode gen.Method, cfg *cfg.GreenConfig) (*bytes.Buffer, *bytes.Buffer, error) {
    89  	outbuf := bytes.NewBuffer(make([]byte, 0, 4096))
    90  	writePkgHeader(outbuf, f.Package)
    91  
    92  	myImports := []string{"fmt"}
    93  	myImports = append(myImports, "github.com/glycerine/greenpack/msgp")
    94  	for _, imp := range f.Imports {
    95  		if imp.Name != nil {
    96  			// have an alias, include it.
    97  			myImports = append(myImports, imp.Name.Name+` `+imp.Path.Value)
    98  		} else {
    99  			myImports = append(myImports, imp.Path.Value)
   100  		}
   101  	}
   102  	dedup := dedupImports(myImports)
   103  	writeImportHeader(outbuf, dedup...)
   104  
   105  	var testbuf *bytes.Buffer
   106  	var testwr io.Writer
   107  	if mode&gen.Test == gen.Test {
   108  		testbuf = bytes.NewBuffer(make([]byte, 0, 4096))
   109  		writePkgHeader(testbuf, f.Package)
   110  		if mode&(gen.Encode|gen.Decode) != 0 {
   111  			writeImportHeader(testbuf, "bytes", "github.com/glycerine/greenpack/msgp", "testing")
   112  		} else {
   113  			writeImportHeader(testbuf, "github.com/glycerine/greenpack/msgp", "testing")
   114  		}
   115  		testwr = testbuf
   116  	}
   117  	return outbuf, testbuf, f.PrintTo(gen.NewPrinter(mode, outbuf, testwr, f.Cfg))
   118  }
   119  
   120  func writePkgHeader(b *bytes.Buffer, name string) {
   121  	b.WriteString("package ")
   122  	b.WriteString(name)
   123  	b.WriteByte('\n')
   124  	b.WriteString("// NOTE: THIS FILE WAS PRODUCED BY THE\n// GREENPACK CODE GENERATION TOOL (github.com/glycerine/greenpack)\n// DO NOT EDIT\n\n")
   125  }
   126  
   127  func writeImportHeader(b *bytes.Buffer, imports ...string) {
   128  	b.WriteString("import (\n")
   129  	for _, im := range imports {
   130  		if im[len(im)-1] == '"' {
   131  			// support aliased imports
   132  			fmt.Fprintf(b, "\t%s\n", im)
   133  		} else {
   134  			fmt.Fprintf(b, "\t%q\n", im)
   135  		}
   136  	}
   137  	b.WriteString(")\n\n")
   138  }
   139  
   140  func fileNameToStructName(file string) string {
   141  	s := strings.TrimSuffix(file, "_gen.go")
   142  	s = strings.Replace(s, "-", "_", -1)
   143  	s = strings.Replace(s, ".", "_", -1)
   144  	s = "File" + strings.ToUpper(s[:1]) + s[1:]
   145  	return s
   146  }