gen.go 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317
  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. // keep track of types with selfer methods
  163. // selferMethods := []string{"CodecEncodeSelf", "CodecDecodeSelf"}
  164. selferEncTyps := make(map[string]bool)
  165. selferDecTyps := make(map[string]bool)
  166. for _, f := range astfiles {
  167. for _, d := range f.Decls {
  168. // if fd, ok := d.(*ast.FuncDecl); ok && fd.Recv != nil && fd.Recv.NumFields() == 1 {
  169. if fd, ok := d.(*ast.FuncDecl); ok && fd.Recv != nil && len(fd.Recv.List) == 1 {
  170. recvType := fd.Recv.List[0].Type
  171. if ptr, ok := recvType.(*ast.StarExpr); ok {
  172. recvType = ptr.X
  173. }
  174. if id, ok := recvType.(*ast.Ident); ok {
  175. switch fd.Name.Name {
  176. case "CodecEncodeSelf":
  177. selferEncTyps[id.Name] = true
  178. case "CodecDecodeSelf":
  179. selferDecTyps[id.Name] = true
  180. }
  181. }
  182. }
  183. }
  184. }
  185. // now find types
  186. for _, f := range astfiles {
  187. for _, d := range f.Decls {
  188. if gd, ok := d.(*ast.GenDecl); ok {
  189. for _, dd := range gd.Specs {
  190. if td, ok := dd.(*ast.TypeSpec); ok {
  191. // if len(td.Name.Name) == 0 || td.Name.Name[0] > 'Z' || td.Name.Name[0] < 'A' {
  192. if len(td.Name.Name) == 0 {
  193. continue
  194. }
  195. // only generate for:
  196. // struct: StructType
  197. // primitives (numbers, bool, string): Ident
  198. // map: MapType
  199. // slice, array: ArrayType
  200. // chan: ChanType
  201. // do not generate:
  202. // FuncType, InterfaceType, StarExpr (ptr), etc
  203. switch td.Type.(type) {
  204. case *ast.StructType, *ast.Ident, *ast.MapType, *ast.ArrayType, *ast.ChanType:
  205. // only add to tv.Types iff
  206. // - it matches per the -r parameter
  207. // - it doesn't match per the -nr parameter
  208. // - it doesn't have any of the Selfer methods in the file
  209. if regexName.FindStringIndex(td.Name.Name) != nil &&
  210. notRegexName.FindStringIndex(td.Name.Name) == nil &&
  211. !selferEncTyps[td.Name.Name] &&
  212. !selferDecTyps[td.Name.Name] {
  213. tv.Types = append(tv.Types, td.Name.Name)
  214. }
  215. }
  216. }
  217. }
  218. }
  219. }
  220. }
  221. if len(tv.Types) == 0 {
  222. return
  223. }
  224. // we cannot use ioutil.TempFile, because we cannot guarantee the file suffix (.go).
  225. // Also, we cannot create file in temp directory,
  226. // because go run will not work (as it needs to see the types here).
  227. // Consequently, create the temp file in the current directory, and remove when done.
  228. // frun, err = ioutil.TempFile("", "codecgen-")
  229. // frunName := filepath.Join(os.TempDir(), "codecgen-"+strconv.FormatInt(time.Now().UnixNano(), 10)+".go")
  230. frunMainName := "codecgen-main-" + tv.RandString + ".generated.go"
  231. frunPkgName := "codecgen-pkg-" + tv.RandString + ".generated.go"
  232. if deleteTempFile {
  233. defer os.Remove(frunMainName)
  234. defer os.Remove(frunPkgName)
  235. }
  236. // var frunMain, frunPkg *os.File
  237. if _, err = gen1(frunMainName, genFrunMainTmpl, &tv); err != nil {
  238. return
  239. }
  240. if _, err = gen1(frunPkgName, genFrunPkgTmpl, &tv); err != nil {
  241. return
  242. }
  243. // remove outfile, so "go run ..." will not think that types in outfile already exist.
  244. os.Remove(outfile)
  245. // execute go run frun
  246. cmd := exec.Command("go", "run", "-tags="+goRunTag, frunMainName) //, frunPkg.Name())
  247. var buf bytes.Buffer
  248. cmd.Stdout = &buf
  249. cmd.Stderr = &buf
  250. if err = cmd.Run(); err != nil {
  251. err = fmt.Errorf("error running 'go run %s': %v, console: %s",
  252. frunMainName, err, buf.Bytes())
  253. return
  254. }
  255. os.Stdout.Write(buf.Bytes())
  256. return
  257. }
  258. func gen1(frunName, tmplStr string, tv interface{}) (frun *os.File, err error) {
  259. os.Remove(frunName)
  260. if frun, err = os.Create(frunName); err != nil {
  261. return
  262. }
  263. defer frun.Close()
  264. t := template.New("")
  265. if t, err = t.Parse(tmplStr); err != nil {
  266. return
  267. }
  268. bw := bufio.NewWriter(frun)
  269. if err = t.Execute(bw, tv); err != nil {
  270. return
  271. }
  272. if err = bw.Flush(); err != nil {
  273. return
  274. }
  275. return
  276. }
  277. func main() {
  278. o := flag.String("o", "", "out file")
  279. c := flag.String("c", genCodecPath, "codec path")
  280. t := flag.String("t", "", "build tag to put in file")
  281. r := flag.String("r", ".*", "regex for type name to match")
  282. nr := flag.String("nr", "^$", "regex for type name to exclude")
  283. rt := flag.String("rt", "", "tags for go run")
  284. st := flag.String("st", "codec,json", "struct tag keys to introspect")
  285. x := flag.Bool("x", false, "keep temp file")
  286. u := flag.Bool("u", false, "Use unsafe, e.g. to avoid unnecessary allocation on []byte->string")
  287. d := flag.Int64("d", 0, "random identifier for use in generated code")
  288. flag.Parse()
  289. if err := Generate(*o, *t, *c, *d, *u, *rt, *st,
  290. regexp.MustCompile(*r), regexp.MustCompile(*nr), !*x, flag.Args()...); err != nil {
  291. fmt.Fprintf(os.Stderr, "codecgen error: %v\n", err)
  292. os.Exit(1)
  293. }
  294. }