main.go 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287
  1. // This program reads all assertion functions from the assert package and
  2. // automatically generates the corersponding requires and forwarded assertions
  3. package main
  4. import (
  5. "bytes"
  6. "flag"
  7. "fmt"
  8. "go/ast"
  9. "go/build"
  10. "go/doc"
  11. "go/importer"
  12. "go/parser"
  13. "go/token"
  14. "go/types"
  15. "io"
  16. "io/ioutil"
  17. "log"
  18. "os"
  19. "path"
  20. "strings"
  21. "text/template"
  22. "github.com/ernesto-jimenez/gogen/imports"
  23. )
  24. var (
  25. pkg = flag.String("assert-path", "github.com/stretchr/testify/assert", "Path to the assert package")
  26. outputPkg = flag.String("output-package", "", "package for the resulting code")
  27. tmplFile = flag.String("template", "", "What file to load the function template from")
  28. out = flag.String("out", "", "What file to write the source code to")
  29. )
  30. func main() {
  31. flag.Parse()
  32. scope, docs, err := parsePackageSource(*pkg)
  33. if err != nil {
  34. log.Fatal(err)
  35. }
  36. importer, funcs, err := analyzeCode(scope, docs)
  37. if err != nil {
  38. log.Fatal(err)
  39. }
  40. if err := generateCode(importer, funcs); err != nil {
  41. log.Fatal(err)
  42. }
  43. }
  44. func generateCode(importer imports.Importer, funcs []testFunc) error {
  45. buff := bytes.NewBuffer(nil)
  46. tmplHead, tmplFunc, err := parseTemplates()
  47. if err != nil {
  48. return err
  49. }
  50. // Generate header
  51. if err := tmplHead.Execute(buff, struct {
  52. Name string
  53. Imports map[string]string
  54. }{
  55. *outputPkg,
  56. importer.Imports(),
  57. }); err != nil {
  58. return err
  59. }
  60. // Generate funcs
  61. for _, fn := range funcs {
  62. buff.Write([]byte("\n\n"))
  63. if err := tmplFunc.Execute(buff, &fn); err != nil {
  64. return err
  65. }
  66. }
  67. // Write file
  68. output, err := outputFile()
  69. if err != nil {
  70. return err
  71. }
  72. defer output.Close()
  73. _, err = io.Copy(output, buff)
  74. return err
  75. }
  76. func parseTemplates() (*template.Template, *template.Template, error) {
  77. tmplHead, err := template.New("header").Parse(headerTemplate)
  78. if err != nil {
  79. return nil, nil, err
  80. }
  81. if *tmplFile != "" {
  82. f, err := ioutil.ReadFile(*tmplFile)
  83. if err != nil {
  84. return nil, nil, err
  85. }
  86. funcTemplate = string(f)
  87. }
  88. tmpl, err := template.New("function").Parse(funcTemplate)
  89. if err != nil {
  90. return nil, nil, err
  91. }
  92. return tmplHead, tmpl, nil
  93. }
  94. func outputFile() (*os.File, error) {
  95. filename := *out
  96. if filename == "-" || (filename == "" && *tmplFile == "") {
  97. return os.Stdout, nil
  98. }
  99. if filename == "" {
  100. filename = strings.TrimSuffix(strings.TrimSuffix(*tmplFile, ".tmpl"), ".go") + ".go"
  101. }
  102. return os.Create(filename)
  103. }
  104. // analyzeCode takes the types scope and the docs and returns the import
  105. // information and information about all the assertion functions.
  106. func analyzeCode(scope *types.Scope, docs *doc.Package) (imports.Importer, []testFunc, error) {
  107. testingT := scope.Lookup("TestingT").Type().Underlying().(*types.Interface)
  108. importer := imports.New(*outputPkg)
  109. var funcs []testFunc
  110. // Go through all the top level functions
  111. for _, fdocs := range docs.Funcs {
  112. // Find the function
  113. obj := scope.Lookup(fdocs.Name)
  114. fn, ok := obj.(*types.Func)
  115. if !ok {
  116. continue
  117. }
  118. // Check function signatuer has at least two arguments
  119. sig := fn.Type().(*types.Signature)
  120. if sig.Params().Len() < 2 {
  121. continue
  122. }
  123. // Check first argument is of type testingT
  124. first, ok := sig.Params().At(0).Type().(*types.Named)
  125. if !ok {
  126. continue
  127. }
  128. firstType, ok := first.Underlying().(*types.Interface)
  129. if !ok {
  130. continue
  131. }
  132. if !types.Implements(firstType, testingT) {
  133. continue
  134. }
  135. funcs = append(funcs, testFunc{*outputPkg, fdocs, fn})
  136. importer.AddImportsFrom(sig.Params())
  137. }
  138. return importer, funcs, nil
  139. }
  140. // parsePackageSource returns the types scope and the package documentation from the pa
  141. func parsePackageSource(pkg string) (*types.Scope, *doc.Package, error) {
  142. pd, err := build.Import(pkg, ".", 0)
  143. if err != nil {
  144. return nil, nil, err
  145. }
  146. fset := token.NewFileSet()
  147. files := make(map[string]*ast.File)
  148. fileList := make([]*ast.File, len(pd.GoFiles))
  149. for i, fname := range pd.GoFiles {
  150. src, err := ioutil.ReadFile(path.Join(pd.SrcRoot, pd.ImportPath, fname))
  151. if err != nil {
  152. return nil, nil, err
  153. }
  154. f, err := parser.ParseFile(fset, fname, src, parser.ParseComments|parser.AllErrors)
  155. if err != nil {
  156. return nil, nil, err
  157. }
  158. files[fname] = f
  159. fileList[i] = f
  160. }
  161. cfg := types.Config{
  162. Importer: importer.Default(),
  163. }
  164. info := types.Info{
  165. Defs: make(map[*ast.Ident]types.Object),
  166. }
  167. tp, err := cfg.Check(pkg, fset, fileList, &info)
  168. if err != nil {
  169. return nil, nil, err
  170. }
  171. scope := tp.Scope()
  172. ap, _ := ast.NewPackage(fset, files, nil, nil)
  173. docs := doc.New(ap, pkg, 0)
  174. return scope, docs, nil
  175. }
  176. type testFunc struct {
  177. CurrentPkg string
  178. DocInfo *doc.Func
  179. TypeInfo *types.Func
  180. }
  181. func (f *testFunc) Qualifier(p *types.Package) string {
  182. if p == nil || p.Name() == f.CurrentPkg {
  183. return ""
  184. }
  185. return p.Name()
  186. }
  187. func (f *testFunc) Params() string {
  188. sig := f.TypeInfo.Type().(*types.Signature)
  189. params := sig.Params()
  190. p := ""
  191. comma := ""
  192. to := params.Len()
  193. var i int
  194. if sig.Variadic() {
  195. to--
  196. }
  197. for i = 1; i < to; i++ {
  198. param := params.At(i)
  199. p += fmt.Sprintf("%s%s %s", comma, param.Name(), types.TypeString(param.Type(), f.Qualifier))
  200. comma = ", "
  201. }
  202. if sig.Variadic() {
  203. param := params.At(params.Len() - 1)
  204. p += fmt.Sprintf("%s%s ...%s", comma, param.Name(), types.TypeString(param.Type().(*types.Slice).Elem(), f.Qualifier))
  205. }
  206. return p
  207. }
  208. func (f *testFunc) ForwardedParams() string {
  209. sig := f.TypeInfo.Type().(*types.Signature)
  210. params := sig.Params()
  211. p := ""
  212. comma := ""
  213. to := params.Len()
  214. var i int
  215. if sig.Variadic() {
  216. to--
  217. }
  218. for i = 1; i < to; i++ {
  219. param := params.At(i)
  220. p += fmt.Sprintf("%s%s", comma, param.Name())
  221. comma = ", "
  222. }
  223. if sig.Variadic() {
  224. param := params.At(params.Len() - 1)
  225. p += fmt.Sprintf("%s%s...", comma, param.Name())
  226. }
  227. return p
  228. }
  229. func (f *testFunc) Comment() string {
  230. return "// " + strings.Replace(strings.TrimSpace(f.DocInfo.Doc), "\n", "\n// ", -1)
  231. }
  232. func (f *testFunc) CommentWithoutT(receiver string) string {
  233. search := fmt.Sprintf("assert.%s(t, ", f.DocInfo.Name)
  234. replace := fmt.Sprintf("%s.%s(", receiver, f.DocInfo.Name)
  235. return strings.Replace(f.Comment(), search, replace, -1)
  236. }
  237. var headerTemplate = `/*
  238. * CODE GENERATED AUTOMATICALLY WITH github.com/stretchr/testify/_codegen
  239. * THIS FILE MUST NOT BE EDITED BY HAND
  240. */
  241. package {{.Name}}
  242. import (
  243. {{range $path, $name := .Imports}}
  244. {{$name}} "{{$path}}"{{end}}
  245. )
  246. `
  247. var funcTemplate = `{{.Comment}}
  248. func (fwd *AssertionsForwarder) {{.DocInfo.Name}}({{.Params}}) bool {
  249. return assert.{{.DocInfo.Name}}({{.ForwardedParams}})
  250. }`