1 // This program reads all assertion functions from the assert package and
2 // automatically generates the corresponding requires and forwarded assertions
27 "github.com/ernesto-jimenez/gogen/imports"
31 pkg = flag.String("assert-path", "github.com/stretchr/testify/assert", "Path to the assert package")
32 includeF = flag.Bool("include-format-funcs", false, "include format functions such as Errorf and Equalf")
33 outputPkg = flag.String("output-package", "", "package for the resulting code")
34 tmplFile = flag.String("template", "", "What file to load the function template from")
35 out = flag.String("out", "", "What file to write the source code to")
41 scope, docs, err := parsePackageSource(*pkg)
46 importer, funcs, err := analyzeCode(scope, docs)
51 if err := generateCode(importer, funcs); err != nil {
56 func generateCode(importer imports.Importer, funcs []testFunc) error {
57 buff := bytes.NewBuffer(nil)
59 tmplHead, tmplFunc, err := parseTemplates()
65 if err := tmplHead.Execute(buff, struct {
67 Imports map[string]string
76 for _, fn := range funcs {
77 buff.Write([]byte("\n\n"))
78 if err := tmplFunc.Execute(buff, &fn); err != nil {
83 code, err := format.Source(buff.Bytes())
89 output, err := outputFile()
94 _, err = io.Copy(output, bytes.NewReader(code))
98 func parseTemplates() (*template.Template, *template.Template, error) {
99 tmplHead, err := template.New("header").Parse(headerTemplate)
104 f, err := ioutil.ReadFile(*tmplFile)
108 funcTemplate = string(f)
110 tmpl, err := template.New("function").Parse(funcTemplate)
114 return tmplHead, tmpl, nil
117 func outputFile() (*os.File, error) {
119 if filename == "-" || (filename == "" && *tmplFile == "") {
120 return os.Stdout, nil
123 filename = strings.TrimSuffix(strings.TrimSuffix(*tmplFile, ".tmpl"), ".go") + ".go"
125 return os.Create(filename)
128 // analyzeCode takes the types scope and the docs and returns the import
129 // information and information about all the assertion functions.
130 func analyzeCode(scope *types.Scope, docs *doc.Package) (imports.Importer, []testFunc, error) {
131 testingT := scope.Lookup("TestingT").Type().Underlying().(*types.Interface)
133 importer := imports.New(*outputPkg)
135 // Go through all the top level functions
136 for _, fdocs := range docs.Funcs {
138 obj := scope.Lookup(fdocs.Name)
140 fn, ok := obj.(*types.Func)
144 // Check function signature has at least two arguments
145 sig := fn.Type().(*types.Signature)
146 if sig.Params().Len() < 2 {
149 // Check first argument is of type testingT
150 first, ok := sig.Params().At(0).Type().(*types.Named)
154 firstType, ok := first.Underlying().(*types.Interface)
158 if !types.Implements(firstType, testingT) {
162 // Skip functions ending with f
163 if strings.HasSuffix(fdocs.Name, "f") && !*includeF {
167 funcs = append(funcs, testFunc{*outputPkg, fdocs, fn})
168 importer.AddImportsFrom(sig.Params())
170 return importer, funcs, nil
173 // parsePackageSource returns the types scope and the package documentation from the package
174 func parsePackageSource(pkg string) (*types.Scope, *doc.Package, error) {
175 pd, err := build.Import(pkg, ".", 0)
180 fset := token.NewFileSet()
181 files := make(map[string]*ast.File)
182 fileList := make([]*ast.File, len(pd.GoFiles))
183 for i, fname := range pd.GoFiles {
184 src, err := ioutil.ReadFile(path.Join(pd.SrcRoot, pd.ImportPath, fname))
188 f, err := parser.ParseFile(fset, fname, src, parser.ParseComments|parser.AllErrors)
197 Importer: importer.Default(),
200 Defs: make(map[*ast.Ident]types.Object),
202 tp, err := cfg.Check(pkg, fset, fileList, &info)
209 ap, _ := ast.NewPackage(fset, files, nil, nil)
210 docs := doc.New(ap, pkg, 0)
212 return scope, docs, nil
215 type testFunc struct {
221 func (f *testFunc) Qualifier(p *types.Package) string {
222 if p == nil || p.Name() == f.CurrentPkg {
228 func (f *testFunc) Params() string {
229 sig := f.TypeInfo.Type().(*types.Signature)
230 params := sig.Params()
239 for i = 1; i < to; i++ {
240 param := params.At(i)
241 p += fmt.Sprintf("%s%s %s", comma, param.Name(), types.TypeString(param.Type(), f.Qualifier))
245 param := params.At(params.Len() - 1)
246 p += fmt.Sprintf("%s%s ...%s", comma, param.Name(), types.TypeString(param.Type().(*types.Slice).Elem(), f.Qualifier))
251 func (f *testFunc) ForwardedParams() string {
252 sig := f.TypeInfo.Type().(*types.Signature)
253 params := sig.Params()
262 for i = 1; i < to; i++ {
263 param := params.At(i)
264 p += fmt.Sprintf("%s%s", comma, param.Name())
268 param := params.At(params.Len() - 1)
269 p += fmt.Sprintf("%s%s...", comma, param.Name())
274 func (f *testFunc) ParamsFormat() string {
275 return strings.Replace(f.Params(), "msgAndArgs", "msg string, args", 1)
278 func (f *testFunc) ForwardedParamsFormat() string {
279 return strings.Replace(f.ForwardedParams(), "msgAndArgs", "append([]interface{}{msg}, args...)", 1)
282 func (f *testFunc) Comment() string {
283 return "// " + strings.Replace(strings.TrimSpace(f.DocInfo.Doc), "\n", "\n// ", -1)
286 func (f *testFunc) CommentFormat() string {
287 search := fmt.Sprintf("%s", f.DocInfo.Name)
288 replace := fmt.Sprintf("%sf", f.DocInfo.Name)
289 comment := strings.Replace(f.Comment(), search, replace, -1)
290 exp := regexp.MustCompile(replace + `\(((\(\)|[^)])+)\)`)
291 return exp.ReplaceAllString(comment, replace+`($1, "error message %s", "formatted")`)
294 func (f *testFunc) CommentWithoutT(receiver string) string {
295 search := fmt.Sprintf("assert.%s(t, ", f.DocInfo.Name)
296 replace := fmt.Sprintf("%s.%s(", receiver, f.DocInfo.Name)
297 return strings.Replace(f.Comment(), search, replace, -1)
300 var headerTemplate = `/*
301 * CODE GENERATED AUTOMATICALLY WITH github.com/stretchr/testify/_codegen
302 * THIS FILE MUST NOT BE EDITED BY HAND
308 {{range $path, $name := .Imports}}
309 {{$name}} "{{$path}}"{{end}}
313 var funcTemplate = `{{.Comment}}
314 func (fwd *AssertionsForwarder) {{.DocInfo.Name}}({{.Params}}) bool {
315 return assert.{{.DocInfo.Name}}({{.ForwardedParams}})