123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285 |
- // Copyright (c) 2012-2015 Ugorji Nwoke. All rights reserved.
- // Use of this source code is governed by a MIT license found in the LICENSE file.
- // codecgen generates codec.Selfer implementations for a set of types.
- package main
- import (
- "bufio"
- "bytes"
- "errors"
- "flag"
- "fmt"
- "go/ast"
- "go/build"
- "go/parser"
- "go/token"
- "math/rand"
- "os"
- "os/exec"
- "path/filepath"
- "regexp"
- "strconv"
- "strings"
- "text/template"
- "time"
- )
- const genCodecPkg = "codec1978" // keep this in sync with codec.genCodecPkg
- const genFrunMainTmpl = `//+build ignore
- package main
- {{ if .Types }}import "{{ .ImportPath }}"{{ end }}
- func main() {
- {{ $.PackageName }}.CodecGenTempWrite{{ .RandString }}()
- }
- `
- // const genFrunPkgTmpl = `//+build codecgen
- const genFrunPkgTmpl = `
- package {{ $.PackageName }}
- import (
- {{ if not .CodecPkgFiles }}{{ .CodecPkgName }} "{{ .CodecImportPath }}"{{ end }}
- "os"
- "reflect"
- "bytes"
- "strings"
- "go/format"
- )
- func CodecGenTempWrite{{ .RandString }}() {
- fout, err := os.Create("{{ .OutFile }}")
- if err != nil {
- panic(err)
- }
- defer fout.Close()
- var out bytes.Buffer
-
- var typs []reflect.Type
- {{ range $index, $element := .Types }}
- var t{{ $index }} {{ . }}
- typs = append(typs, reflect.TypeOf(t{{ $index }}))
- {{ end }}
- {{ if not .CodecPkgFiles }}{{ .CodecPkgName }}.{{ end }}Gen(&out, "{{ .BuildTag }}", "{{ .PackageName }}", "{{ .RandString }}", {{ .UseUnsafe }}, {{ if not .CodecPkgFiles }}{{ .CodecPkgName }}.{{ end }}NewTypeInfos(strings.Split("{{ .StructTags }}", ",")), typs...)
- bout, err := format.Source(out.Bytes())
- if err != nil {
- fout.Write(out.Bytes())
- panic(err)
- }
- fout.Write(bout)
- }
- `
- // Generate is given a list of *.go files to parse, and an output file (fout).
- //
- // It finds all types T in the files, and it creates 2 tmp files (frun).
- // - main package file passed to 'go run'
- // - package level file which calls *genRunner.Selfer to write Selfer impls for each T.
- // We use a package level file so that it can reference unexported types in the package being worked on.
- // Tool then executes: "go run __frun__" which creates fout.
- // fout contains Codec(En|De)codeSelf implementations for every type T.
- //
- func Generate(outfile, buildTag, codecPkgPath string, uid int64, useUnsafe bool, goRunTag string,
- st string, regexName *regexp.Regexp, notRegexName *regexp.Regexp, deleteTempFile bool, infiles ...string) (err error) {
- // For each file, grab AST, find each type, and write a call to it.
- if len(infiles) == 0 {
- return
- }
- if outfile == "" || codecPkgPath == "" {
- err = errors.New("outfile and codec package path cannot be blank")
- return
- }
- if uid < 0 {
- uid = -uid
- }
- if uid == 0 {
- rr := rand.New(rand.NewSource(time.Now().UnixNano()))
- uid = 101 + rr.Int63n(9777)
- }
- // We have to parse dir for package, before opening the temp file for writing (else ImportDir fails).
- // Also, ImportDir(...) must take an absolute path.
- lastdir := filepath.Dir(outfile)
- absdir, err := filepath.Abs(lastdir)
- if err != nil {
- return
- }
- pkg, err := build.Default.ImportDir(absdir, build.AllowBinary)
- if err != nil {
- return
- }
- type tmplT struct {
- CodecPkgName string
- CodecImportPath string
- ImportPath string
- OutFile string
- PackageName string
- RandString string
- BuildTag string
- StructTags string
- Types []string
- CodecPkgFiles bool
- UseUnsafe bool
- }
- tv := tmplT{
- CodecPkgName: genCodecPkg,
- OutFile: outfile,
- CodecImportPath: codecPkgPath,
- BuildTag: buildTag,
- UseUnsafe: useUnsafe,
- RandString: strconv.FormatInt(uid, 10),
- StructTags: st,
- }
- tv.ImportPath = pkg.ImportPath
- if tv.ImportPath == tv.CodecImportPath {
- tv.CodecPkgFiles = true
- tv.CodecPkgName = "codec"
- } else {
- // HACK: always handle vendoring. It should be typically on in go 1.6, 1.7
- s := tv.ImportPath
- const vendorStart = "vendor/"
- const vendorInline = "/vendor/"
- if i := strings.LastIndex(s, vendorInline); i >= 0 {
- tv.ImportPath = s[i+len(vendorInline):]
- } else if strings.HasPrefix(s, vendorStart) {
- tv.ImportPath = s[len(vendorStart):]
- }
- }
- astfiles := make([]*ast.File, len(infiles))
- for i, infile := range infiles {
- if filepath.Dir(infile) != lastdir {
- err = errors.New("in files must all be in same directory as outfile")
- return
- }
- fset := token.NewFileSet()
- astfiles[i], err = parser.ParseFile(fset, infile, nil, 0)
- if err != nil {
- return
- }
- if i == 0 {
- tv.PackageName = astfiles[i].Name.Name
- if tv.PackageName == "main" {
- // codecgen cannot be run on types in the 'main' package.
- // A temporary 'main' package must be created, and should reference the fully built
- // package containing the types.
- // Also, the temporary main package will conflict with the main package which already has a main method.
- err = errors.New("codecgen cannot be run on types in the 'main' package")
- return
- }
- }
- }
- for _, f := range astfiles {
- for _, d := range f.Decls {
- if gd, ok := d.(*ast.GenDecl); ok {
- for _, dd := range gd.Specs {
- if td, ok := dd.(*ast.TypeSpec); ok {
- // if len(td.Name.Name) == 0 || td.Name.Name[0] > 'Z' || td.Name.Name[0] < 'A' {
- if len(td.Name.Name) == 0 {
- continue
- }
- // only generate for:
- // struct: StructType
- // primitives (numbers, bool, string): Ident
- // map: MapType
- // slice, array: ArrayType
- // chan: ChanType
- // do not generate:
- // FuncType, InterfaceType, StarExpr (ptr), etc
- switch td.Type.(type) {
- case *ast.StructType, *ast.Ident, *ast.MapType, *ast.ArrayType, *ast.ChanType:
- if regexName.FindStringIndex(td.Name.Name) != nil && notRegexName.FindStringIndex(td.Name.Name) == nil {
- tv.Types = append(tv.Types, td.Name.Name)
- }
- }
- }
- }
- }
- }
- }
- if len(tv.Types) == 0 {
- return
- }
- // we cannot use ioutil.TempFile, because we cannot guarantee the file suffix (.go).
- // Also, we cannot create file in temp directory,
- // because go run will not work (as it needs to see the types here).
- // Consequently, create the temp file in the current directory, and remove when done.
- // frun, err = ioutil.TempFile("", "codecgen-")
- // frunName := filepath.Join(os.TempDir(), "codecgen-"+strconv.FormatInt(time.Now().UnixNano(), 10)+".go")
- frunMainName := "codecgen-main-" + tv.RandString + ".generated.go"
- frunPkgName := "codecgen-pkg-" + tv.RandString + ".generated.go"
- if deleteTempFile {
- defer os.Remove(frunMainName)
- defer os.Remove(frunPkgName)
- }
- // var frunMain, frunPkg *os.File
- if _, err = gen1(frunMainName, genFrunMainTmpl, &tv); err != nil {
- return
- }
- if _, err = gen1(frunPkgName, genFrunPkgTmpl, &tv); err != nil {
- return
- }
- // remove outfile, so "go run ..." will not think that types in outfile already exist.
- os.Remove(outfile)
- // execute go run frun
- cmd := exec.Command("go", "run", "-tags="+goRunTag, frunMainName) //, frunPkg.Name())
- var buf bytes.Buffer
- cmd.Stdout = &buf
- cmd.Stderr = &buf
- if err = cmd.Run(); err != nil {
- err = fmt.Errorf("error running 'go run %s': %v, console: %s",
- frunMainName, err, buf.Bytes())
- return
- }
- os.Stdout.Write(buf.Bytes())
- return
- }
- func gen1(frunName, tmplStr string, tv interface{}) (frun *os.File, err error) {
- os.Remove(frunName)
- if frun, err = os.Create(frunName); err != nil {
- return
- }
- defer frun.Close()
- t := template.New("")
- if t, err = t.Parse(tmplStr); err != nil {
- return
- }
- bw := bufio.NewWriter(frun)
- if err = t.Execute(bw, tv); err != nil {
- return
- }
- if err = bw.Flush(); err != nil {
- return
- }
- return
- }
- func main() {
- o := flag.String("o", "", "out file")
- c := flag.String("c", genCodecPath, "codec path")
- t := flag.String("t", "", "build tag to put in file")
- r := flag.String("r", ".*", "regex for type name to match")
- nr := flag.String("nr", "^$", "regex for type name to exclude")
- rt := flag.String("rt", "", "tags for go run")
- st := flag.String("st", "codec,json", "struct tag keys to introspect")
- x := flag.Bool("x", false, "keep temp file")
- u := flag.Bool("u", false, "Use unsafe, e.g. to avoid unnecessary allocation on []byte->string")
- d := flag.Int64("d", 0, "random identifier for use in generated code")
- flag.Parse()
- if err := Generate(*o, *t, *c, *d, *u, *rt, *st,
- regexp.MustCompile(*r), regexp.MustCompile(*nr), !*x, flag.Args()...); err != nil {
- fmt.Fprintf(os.Stderr, "codecgen error: %v\n", err)
- os.Exit(1)
- }
- }
|