interfacegen.go

package main

import (
	"bytes"
	"context"
	"fmt"
	"go/ast"
	"go/token"
	"go/types"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime/debug"
	"strings"

	"github.com/spf13/cobra"
	"golang.org/x/tools/go/packages"
	"golang.org/x/tools/go/types/typeutil"
	"golang.org/x/tools/imports"
)

var (
	skip = regexp.MustCompile(`//\s*interfacegen:skip`)
)

func main() {
	var types []string
	app := application{}

	cmd := &cobra.Command{
		Use:     "interfacegen",
		Short:   "interfacegen generates interfaces from concrete types",
		Version: buildVersionString(),
		Run: func(cmd *cobra.Command, args []string) {
			for _, t := range types {
				components := strings.Split(t, ":")
				wl := typeWhitelist{in: components[0]}
				if len(components) > 1 {
					wl.out = components[1]
				}

				app.TypeWhitelist = append(app.TypeWhitelist, wl)
			}

			err := app.Run(context.Background())
			if err != nil {
				log.Fatal(err)
			}
		},
	}

	cmd.Flags().StringVarP(&app.Output, "output", "o", "-", "output file for generated code")
	cmd.Flags().StringVarP(&app.SrcPackage, "src", "s", "", "source package path")
	cmd.Flags().StringVarP(&app.DstPackage, "dst", "d", "interfaces", "destination package name")
	cmd.Flags().StringVarP(&app.Comment, "comment", "c", "Code generated by interfacegen. DO NOT EDIT.", "comment for generated packages")
	cmd.Flags().BoolVar(&app.IncludeDocs, "doc", true, "whether to include method comments")
	cmd.Flags().BoolVarP(&app.IncludePrivate, "private", "p", false, "include private methods")
	cmd.Flags().StringSliceVarP(&types, "types", "t", nil, "whitelist of types to include in concreteName:interfaceName format (can repeat)")
	cmd.Flags().BoolVarP(&app.IncludeAllPackages, "all", "a", true, "include all types")

	err := cmd.MarkFlagRequired("src")
	if err != nil {
		log.Fatal(err)
	}

	err = cmd.Execute()
	if err != nil {
		log.Fatal(err)
	}
}

type typeWhitelist struct {
	in  string
	out string
}

type application struct {
	IncludePrivate     bool
	SrcPackage         string
	DstPackage         string
	TypeWhitelist      []typeWhitelist
	IncludeAllPackages bool
	Output             string
	Comment            string
	IncludeDocs        bool
	BuildFlags         []string `flag:"buildflag" help:"pass argument to underlying build system (may be repeated)"`
}

// Run takes the args after flag processing and performs the specified query.
func (app *application) Run(ctx context.Context) error {
	// Load, parse, and type-check the packages named on the command line.
	mode := (packages.NeedName | packages.NeedFiles | packages.NeedImports |
		packages.NeedTypes | packages.NeedTypesSizes | packages.NeedSyntax | packages.NeedTypesInfo | packages.NeedModule)

	cfg := &packages.Config{
		Mode:       mode,
		Tests:      false,
		BuildFlags: app.BuildFlags,
	}

	lpkgs, err := packages.Load(cfg, app.SrcPackage)
	if err != nil {
		return err
	}

	return app.parse(lpkgs)
}

func (app *application) parse(lpkgs []*packages.Package) error {
	imports := map[string]bool{}

	var interfaceDefs []interfaceDef

	for _, lpkg := range lpkgs {
		comments := populateDocs(lpkg)

		var target *types.Package

		if app.DstPackage != "" {
			target = types.NewPackage(app.DstPackage, app.DstPackage)
		} else {
			target = lpkg.Types
		}

		qualifier := func(other *types.Package) string {
			if target == other {
				return ""
			}

			return other.Name()
		}

		scope := lpkg.Types.Scope()
		for _, name := range scope.Names() {
			obj := scope.Lookup(name)
			if !obj.Exported() && !app.IncludePrivate {
				continue
			}

			if _, ok := obj.(*types.TypeName); ok {
				name := obj.Name()

				if len(app.TypeWhitelist) > 0 {
					var found bool

					for _, t := range app.TypeWhitelist {
						if t.in == obj.Name() {
							found = true
							if t.out != "" {
								name = t.out
							}

							break
						}
					}

					if !found && !app.IncludeAllPackages {
						continue
					}
				}

				var methods []methodDef

				for _, meth := range typeutil.IntuitiveMethodSet(obj.Type(), nil) {
					if !meth.Obj().Exported() && !app.IncludePrivate {
						continue
					}
					b := bytes.NewBuffer(nil)
					sig := meth.Obj().Type().(*types.Signature)
					types.WriteSignature(b, sig, qualifier)

					slug := strings.Join([]string{lpkg.ID, obj.Name(), meth.Obj().Name()}, ".")
					if shouldSkip(comments[slug]) {
						continue
					}

					doc := comments[slug]
					if !app.IncludeDocs {
						doc = nil
					}

					methods = append(methods, methodDef{
						Method: meth.Obj().Name() + b.String(),
						Slug:   slug,
						Doc:    doc,
					})
				}

				if len(methods) == 0 {
					continue
				}

				slug := strings.Join([]string{lpkg.ID, name}, ".")
				if shouldSkip(comments[slug]) {
					continue
				}

				doc := comments[slug]
				if !app.IncludeDocs {
					doc = nil
				}

				interfaceDefs = append(interfaceDefs, interfaceDef{
					Name:    name,
					Doc:     doc,
					Methods: methods,
				})
			}
		}

		if len(interfaceDefs) == 0 {
			continue
		}

		for _, imp := range lpkg.Imports {
			imports[imp.ID] = true
		}

		if target != lpkg.Types {
			imports[lpkg.ID] = true
		}
	}

	if len(interfaceDefs) == 0 {
		return fmt.Errorf("package %s contained no matching interfaces", app.SrcPackage)
	}

	var distinctImports []string
	for path := range imports {
		distinctImports = append(distinctImports, path)
	}

	code, err := app.generatePackage(distinctImports, interfaceDefs)
	if err != nil {
		return err
	}

	var formatted []byte

	// Continue formatting the output until it stabilizes, as goimports
	// can't currently remove duplicate unused imports in one shot.
	for {
		formatted, err = format([]byte(code))
		if err != nil {
			return err
		}

		if string(formatted) == code {
			break
		}

		code = string(formatted)
	}

	switch app.Output {
	case "", "-":
		fmt.Println(string(formatted))
	default:
		return app.writeFile(formatted)
	}

	return nil
}

func (app *application) writeFile(data []byte) (err error) {
	err = os.MkdirAll(filepath.Dir(app.Output), 0755)
	if err != nil {
		return err
	}

	w, err := os.OpenFile(app.Output, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer w.Close()

	_, err = w.Write(data)
	if err != nil {
		return err
	}

	return nil
}

func format(code []byte) (formatted []byte, err error) {
	opts := &imports.Options{
		TabIndent: true,
		TabWidth:  2,
		Fragment:  true,
		Comments:  true,
	}
	return imports.Process("", code, opts)
}

func populateDocs(p *packages.Package) map[string][]string {
	ret := map[string][]string{}

	for _, f := range p.Syntax {
		for _, decl := range f.Decls {
			switch v := decl.(type) {
			case *ast.FuncDecl:
				if v.Recv == nil || v.Doc == nil {
					continue
				}

				se, ok := v.Recv.List[0].Type.(*ast.StarExpr)
				if !ok {
					continue
				}

				ident, ok := se.X.(*ast.Ident)
				if !ok {
					continue
				}

				var docs []string

				for _, d := range v.Doc.List {
					docs = append(docs, d.Text)
				}

				if len(docs) == 0 {
					continue
				}

				slug := strings.Join([]string{p.ID, ident.Name, v.Name.String()}, ".")
				ret[slug] = docs
			case *ast.GenDecl:
				if v.Tok != token.TYPE {
					continue
				}

				spec := v.Specs[0].(*ast.TypeSpec)

				var docs []string

				if v.Doc == nil {
					continue
				}

				for _, d := range v.Doc.List {
					docs = append(docs, d.Text)
				}

				if len(docs) == 0 {
					continue
				}

				slug := strings.Join([]string{p.ID, spec.Name.String()}, ".")
				ret[slug] = docs
			}
		}
	}

	return ret
}

func shouldSkip(comments []string) bool {
	for _, comment := range comments {
		m := skip.MatchString(comment)
		if m {
			return true
		}
	}

	return false
}

func buildVersionString() string {
	var s string
	info, ok := debug.ReadBuildInfo()
	if !ok {
		return "\nunknown"
	}

	if info.Main.Path != "" {
		s += "\nmodule: " + info.Main.Path
	}

	if info.GoVersion != "" {
		s += "\ncompiler: " + info.GoVersion
	}

	if info.Main.Version != "" {
		s += "\nsemver: " + info.Main.Version
	}

	for _, buildSetting := range info.Settings {
		if buildSetting.Value == "" {
			continue
		}
		switch buildSetting.Key {
		case "vcs.revision":
			s += "\nvcs.revision: " + buildSetting.Value
		case "vcs.time":
			s += "\nvcs.revision.time: " + buildSetting.Value
		}
	}

	return s
}