gostaticanalysis / codegen

Library: Code generator

codegen

codegen is a library for code generators with analysis package.

package mockgen

import (
	"bytes"
	"errors"
	"fmt"
	"go/format"
	"os"

	"github.com/gostaticanalysis/codegen"
	"github.com/gostaticanalysis/knife"
)

var (
	flagTypeName string
	flagOutput   string
)

func init() {
	Generator.Flags.StringVar(&flagTypeName, "type", "", "type name of interface")
	Generator.Flags.StringVar(&flagOutput, "o", "", "output file name")
}

var Generator = &codegen.Generator{
	Name: "mockgen",
	Doc:  "mockgen generates a mock of interface",
	Run: func(pass *codegen.Pass) error {
		if flagTypeName == "" {
			return errors.New("type must be specified")
		}

		td := &knife.TempalteData{
			Fset:      pass.Fset,
			Files:     pass.Files,
			TypesInfo: pass.TypesInfo,
			Pkg:       pass.Pkg,
			Extra: map[string]interface{}{
				"type": flagTypeName,
			},
		}
		t, err := knife.NewTemplate(td).Parse(tmpl)
		if err != nil {
			return err
		}

		var buf bytes.Buffer
		if err := t.Execute(&buf, knife.NewPackage(pass.Pkg)); err != nil {
			return err
		}

		src, err := format.Source(buf.Bytes())
		if err != nil {
			return err
		}

		if flagOutput == "" {
			pass.Print(string(src))
			return nil
		}

		f, err := os.Create(flagOutput)
		if err != nil {
			return err
		}

		fmt.Fprint(f, string(src))

		if err := f.Close(); err != nil {
			return err
		}

		return nil
	},
}

var tmpl = `{{with index .Types (data "type")}}{{if interface .}}
// Code generated by mockgen; DO NOT EDIT.
package {{(pkg).Name}}
type Mock{{data "type"}} struct {
{{- range $n, $f := methods .}}
	{{$n}}Func {{$f.Signature}}
{{- end}}
}
{{range $n, $f := methods .}}
func (m *Mock{{data "type"}}) {{$n}}({{range $f.Signature.Params}}
	{{- .Name}} {{.Type}},
{{- end}}) ({{range $f.Signature.Results}}
	{{- .Name}} {{.Type}},
{{- end}}) {
	{{if $f.Signature.Results}}return {{end}}m.{{$n}}Func({{range $f.Signature.Params}}
		{{- .Name}},
	{{- end}})
}
{{end}}
{{end}}
{{end}}
`