gen.go 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285
  1. // Copyright (c) 2012-2015 Ugorji Nwoke. All rights reserved.
  2. // Use of this source code is governed by a MIT license found in the LICENSE file.
  3. // codecgen generates codec.Selfer implementations for a set of types.
  4. package main
  5. import (
  6. "bufio"
  7. "bytes"
  8. "errors"
  9. "flag"
  10. "fmt"
  11. "go/ast"
  12. "go/build"
  13. "go/parser"
  14. "go/token"
  15. "math/rand"
  16. "os"
  17. "os/exec"
  18. "path/filepath"
  19. "regexp"
  20. "strconv"
  21. "strings"
  22. "text/template"
  23. "time"
  24. )
  25. const genCodecPkg = "codec1978" // keep this in sync with codec.genCodecPkg
  26. const genFrunMainTmpl = `//+build ignore
  27. package main
  28. {{ if .Types }}import "{{ .ImportPath }}"{{ end }}
  29. func main() {
  30. {{ $.PackageName }}.CodecGenTempWrite{{ .RandString }}()
  31. }
  32. `
  33. // const genFrunPkgTmpl = `//+build codecgen
  34. const genFrunPkgTmpl = `
  35. package {{ $.PackageName }}
  36. import (
  37. {{ if not .CodecPkgFiles }}{{ .CodecPkgName }} "{{ .CodecImportPath }}"{{ end }}
  38. "os"
  39. "reflect"
  40. "bytes"
  41. "strings"
  42. "go/format"
  43. )
  44. func CodecGenTempWrite{{ .RandString }}() {
  45. fout, err := os.Create("{{ .OutFile }}")
  46. if err != nil {
  47. panic(err)
  48. }
  49. defer fout.Close()
  50. var out bytes.Buffer
  51. var typs []reflect.Type
  52. {{ range $index, $element := .Types }}
  53. var t{{ $index }} {{ . }}
  54. typs = append(typs, reflect.TypeOf(t{{ $index }}))
  55. {{ end }}
  56. {{ if not .CodecPkgFiles }}{{ .CodecPkgName }}.{{ end }}Gen(&out, "{{ .BuildTag }}", "{{ .PackageName }}", "{{ .RandString }}", {{ .UseUnsafe }}, {{ if not .CodecPkgFiles }}{{ .CodecPkgName }}.{{ end }}NewTypeInfos(strings.Split("{{ .StructTags }}", ",")), typs...)
  57. bout, err := format.Source(out.Bytes())
  58. if err != nil {
  59. fout.Write(out.Bytes())
  60. panic(err)
  61. }
  62. fout.Write(bout)
  63. }
  64. `
  65. // Generate is given a list of *.go files to parse, and an output file (fout).
  66. //
  67. // It finds all types T in the files, and it creates 2 tmp files (frun).
  68. // - main package file passed to 'go run'
  69. // - package level file which calls *genRunner.Selfer to write Selfer impls for each T.
  70. // We use a package level file so that it can reference unexported types in the package being worked on.
  71. // Tool then executes: "go run __frun__" which creates fout.
  72. // fout contains Codec(En|De)codeSelf implementations for every type T.
  73. //
  74. func Generate(outfile, buildTag, codecPkgPath string, uid int64, useUnsafe bool, goRunTag string,
  75. st string, regexName *regexp.Regexp, notRegexName *regexp.Regexp, deleteTempFile bool, infiles ...string) (err error) {
  76. // For each file, grab AST, find each type, and write a call to it.
  77. if len(infiles) == 0 {
  78. return
  79. }
  80. if outfile == "" || codecPkgPath == "" {
  81. err = errors.New("outfile and codec package path cannot be blank")
  82. return
  83. }
  84. if uid < 0 {
  85. uid = -uid
  86. }
  87. if uid == 0 {
  88. rr := rand.New(rand.NewSource(time.Now().UnixNano()))
  89. uid = 101 + rr.Int63n(9777)
  90. }
  91. // We have to parse dir for package, before opening the temp file for writing (else ImportDir fails).
  92. // Also, ImportDir(...) must take an absolute path.
  93. lastdir := filepath.Dir(outfile)
  94. absdir, err := filepath.Abs(lastdir)
  95. if err != nil {
  96. return
  97. }
  98. pkg, err := build.Default.ImportDir(absdir, build.AllowBinary)
  99. if err != nil {
  100. return
  101. }
  102. type tmplT struct {
  103. CodecPkgName string
  104. CodecImportPath string
  105. ImportPath string
  106. OutFile string
  107. PackageName string
  108. RandString string
  109. BuildTag string
  110. StructTags string
  111. Types []string
  112. CodecPkgFiles bool
  113. UseUnsafe bool
  114. }
  115. tv := tmplT{
  116. CodecPkgName: genCodecPkg,
  117. OutFile: outfile,
  118. CodecImportPath: codecPkgPath,
  119. BuildTag: buildTag,
  120. UseUnsafe: useUnsafe,
  121. RandString: strconv.FormatInt(uid, 10),
  122. StructTags: st,
  123. }
  124. tv.ImportPath = pkg.ImportPath
  125. if tv.ImportPath == tv.CodecImportPath {
  126. tv.CodecPkgFiles = true
  127. tv.CodecPkgName = "codec"
  128. } else {
  129. // HACK: always handle vendoring. It should be typically on in go 1.6, 1.7
  130. s := tv.ImportPath
  131. const vendorStart = "vendor/"
  132. const vendorInline = "/vendor/"
  133. if i := strings.LastIndex(s, vendorInline); i >= 0 {
  134. tv.ImportPath = s[i+len(vendorInline):]
  135. } else if strings.HasPrefix(s, vendorStart) {
  136. tv.ImportPath = s[len(vendorStart):]
  137. }
  138. }
  139. astfiles := make([]*ast.File, len(infiles))
  140. for i, infile := range infiles {
  141. if filepath.Dir(infile) != lastdir {
  142. err = errors.New("in files must all be in same directory as outfile")
  143. return
  144. }
  145. fset := token.NewFileSet()
  146. astfiles[i], err = parser.ParseFile(fset, infile, nil, 0)
  147. if err != nil {
  148. return
  149. }
  150. if i == 0 {
  151. tv.PackageName = astfiles[i].Name.Name
  152. if tv.PackageName == "main" {
  153. // codecgen cannot be run on types in the 'main' package.
  154. // A temporary 'main' package must be created, and should reference the fully built
  155. // package containing the types.
  156. // Also, the temporary main package will conflict with the main package which already has a main method.
  157. err = errors.New("codecgen cannot be run on types in the 'main' package")
  158. return
  159. }
  160. }
  161. }
  162. for _, f := range astfiles {
  163. for _, d := range f.Decls {
  164. if gd, ok := d.(*ast.GenDecl); ok {
  165. for _, dd := range gd.Specs {
  166. if td, ok := dd.(*ast.TypeSpec); ok {
  167. // if len(td.Name.Name) == 0 || td.Name.Name[0] > 'Z' || td.Name.Name[0] < 'A' {
  168. if len(td.Name.Name) == 0 {
  169. continue
  170. }
  171. // only generate for:
  172. // struct: StructType
  173. // primitives (numbers, bool, string): Ident
  174. // map: MapType
  175. // slice, array: ArrayType
  176. // chan: ChanType
  177. // do not generate:
  178. // FuncType, InterfaceType, StarExpr (ptr), etc
  179. switch td.Type.(type) {
  180. case *ast.StructType, *ast.Ident, *ast.MapType, *ast.ArrayType, *ast.ChanType:
  181. if regexName.FindStringIndex(td.Name.Name) != nil && notRegexName.FindStringIndex(td.Name.Name) == nil {
  182. tv.Types = append(tv.Types, td.Name.Name)
  183. }
  184. }
  185. }
  186. }
  187. }
  188. }
  189. }
  190. if len(tv.Types) == 0 {
  191. return
  192. }
  193. // we cannot use ioutil.TempFile, because we cannot guarantee the file suffix (.go).
  194. // Also, we cannot create file in temp directory,
  195. // because go run will not work (as it needs to see the types here).
  196. // Consequently, create the temp file in the current directory, and remove when done.
  197. // frun, err = ioutil.TempFile("", "codecgen-")
  198. // frunName := filepath.Join(os.TempDir(), "codecgen-"+strconv.FormatInt(time.Now().UnixNano(), 10)+".go")
  199. frunMainName := "codecgen-main-" + tv.RandString + ".generated.go"
  200. frunPkgName := "codecgen-pkg-" + tv.RandString + ".generated.go"
  201. if deleteTempFile {
  202. defer os.Remove(frunMainName)
  203. defer os.Remove(frunPkgName)
  204. }
  205. // var frunMain, frunPkg *os.File
  206. if _, err = gen1(frunMainName, genFrunMainTmpl, &tv); err != nil {
  207. return
  208. }
  209. if _, err = gen1(frunPkgName, genFrunPkgTmpl, &tv); err != nil {
  210. return
  211. }
  212. // remove outfile, so "go run ..." will not think that types in outfile already exist.
  213. os.Remove(outfile)
  214. // execute go run frun
  215. cmd := exec.Command("go", "run", "-tags="+goRunTag, frunMainName) //, frunPkg.Name())
  216. var buf bytes.Buffer
  217. cmd.Stdout = &buf
  218. cmd.Stderr = &buf
  219. if err = cmd.Run(); err != nil {
  220. err = fmt.Errorf("error running 'go run %s': %v, console: %s",
  221. frunMainName, err, buf.Bytes())
  222. return
  223. }
  224. os.Stdout.Write(buf.Bytes())
  225. return
  226. }
  227. func gen1(frunName, tmplStr string, tv interface{}) (frun *os.File, err error) {
  228. os.Remove(frunName)
  229. if frun, err = os.Create(frunName); err != nil {
  230. return
  231. }
  232. defer frun.Close()
  233. t := template.New("")
  234. if t, err = t.Parse(tmplStr); err != nil {
  235. return
  236. }
  237. bw := bufio.NewWriter(frun)
  238. if err = t.Execute(bw, tv); err != nil {
  239. return
  240. }
  241. if err = bw.Flush(); err != nil {
  242. return
  243. }
  244. return
  245. }
  246. func main() {
  247. o := flag.String("o", "", "out file")
  248. c := flag.String("c", genCodecPath, "codec path")
  249. t := flag.String("t", "", "build tag to put in file")
  250. r := flag.String("r", ".*", "regex for type name to match")
  251. nr := flag.String("nr", "^$", "regex for type name to exclude")
  252. rt := flag.String("rt", "", "tags for go run")
  253. st := flag.String("st", "codec,json", "struct tag keys to introspect")
  254. x := flag.Bool("x", false, "keep temp file")
  255. u := flag.Bool("u", false, "Use unsafe, e.g. to avoid unnecessary allocation on []byte->string")
  256. d := flag.Int64("d", 0, "random identifier for use in generated code")
  257. flag.Parse()
  258. if err := Generate(*o, *t, *c, *d, *u, *rt, *st,
  259. regexp.MustCompile(*r), regexp.MustCompile(*nr), !*x, flag.Args()...); err != nil {
  260. fmt.Fprintf(os.Stderr, "codecgen error: %v\n", err)
  261. os.Exit(1)
  262. }
  263. }