generator.go

package main

import (
	"bytes"
	"fmt"
	"go/format"
	"go/token"
	"log"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/ssoor/implgen/model"
)

type generator struct {
	buf                       bytes.Buffer
	head                      bool
	dstFileName               string
	indent                    string
	mockNames                 map[string]string // may be empty
	mockInterfaces            map[string]bool   // may be empty
	mockInterfaceRegexps      []*regexp.Regexp
	filename                  string // may be empty
	srcPackage, srcInterfaces string // may be empty
	copyrightHeader           string

	packageMap map[string]string // map from import path to package name
}

func (g *generator) p(format string, args ...interface{}) {
	fmt.Fprintf(&g.buf, g.indent+format+"\n", args...)
}
func (g *generator) pf(format string, args ...interface{}) {
	fmt.Fprintf(&g.buf, g.indent+format, args...)
}

func (g *generator) in() {
	g.indent += "\t"
}

func (g *generator) out() {
	if len(g.indent) > 0 {
		g.indent = g.indent[0 : len(g.indent)-1]
	}
}

func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPackagePath string) error {
	dstPkg, err := sourceMode(g.dstFileName)

	g.generatePackageMap(pkg, outputPkgName, outputPackagePath)

	if err != nil {
		g.head = true
		g.generateHead(pkg, outputPkgName, outputPackagePath)
	} else {
	}

	return g.generate(pkg, dstPkg, outputPkgName, outputPackagePath)
}

func (g *generator) generatePackageMap(pkg *model.Package, outputPkgName string, outputPackagePath string) {
	// Get all required imports, and generate unique names for them all.
	im := pkg.Imports()

	// Only import reflect if it's used. We only use reflect in mocked methods
	// so only import if any of the mocked interfaces have methods.
	for _, intf := range pkg.Interfaces {
		if len(intf.Methods) > 0 {
			break
		}
	}

	// Sort keys to make import alias generation predictable
	sortedPaths := make([]string, len(im))
	x := 0
	for pth := range im {
		sortedPaths[x] = pth
		x++
	}
	sort.Strings(sortedPaths)

	packagesName := createPackageMap(sortedPaths)

	g.packageMap = make(map[string]string, len(im))
	localNames := make(map[string]bool, len(im))
	for _, pth := range sortedPaths {
		base, ok := packagesName[pth]
		if !ok {
			base = sanitize(path.Base(pth))
		}

		// Local names for an imported package can usually be the basename of the import path.
		// A couple of situations don't permit that, such as duplicate local names
		// (e.g. importing "html/template" and "text/template"), or where the basename is
		// a keyword (e.g. "foo/case").
		// try base0, base1, ...
		pkgName := base
		i := 0
		for localNames[pkgName] || token.Lookup(pkgName).IsKeyword() {
			pkgName = base + strconv.Itoa(i)
			i++
		}

		// Avoid importing package if source pkg == output pkg
		if pth == pkg.PkgPath && outputPkgName == pkg.Name {
			continue
		}

		g.packageMap[pth] = pkgName
		localNames[pkgName] = true
	}
}
func (g *generator) generateHead(pkg *model.Package, outputPkgName string, outputPackagePath string) {
	if outputPkgName != pkg.Name && *selfPackage == "" {
		// reset outputPackagePath if it's not passed in through -self_package
		outputPackagePath = ""
	}

	if g.copyrightHeader != "" {
		lines := strings.Split(g.copyrightHeader, "\n")
		for _, line := range lines {
			g.p("// %s", line)
		}
		g.p("")
	}

	g.p("// Code generated by ImplGen.")
	if g.filename != "" {
		g.p("// Source: %v", g.filename)
	} else {
		g.p("// Source: %v (interfaces: %v)", g.srcPackage, g.srcInterfaces)
	}
	g.p("")

	if *writePkgComment {
		g.p("// Package %v is a generated ImplGen package.", outputPkgName)
	}
	g.p("package %v", outputPkgName)
	g.p("")
	g.p("import (")
	g.in()
	for pkgPath, pkgName := range g.packageMap {
		if pkgPath == outputPackagePath {
			continue
		}
		g.p("%v %q", pkgName, pkgPath)
	}
	for _, pkgPath := range pkg.DotImports {
		g.p(". %q", pkgPath)
	}
	g.out()
	g.p(")")
}

func (g *generator) generate(pkg *model.Package, dstPkg *model.Package, outputPkgName string, outputPackagePath string) error {
	namesMap := make(map[string]*model.Struct)
	if dstPkg != nil {
		for _, sn := range dstPkg.StructNames {
			namesMap[sn.Name] = sn
		}
	}

	for _, intf := range pkg.Interfaces {
		if !g.mockInterfaces[intf.Name] {
			match := false
			for _, v := range g.mockInterfaceRegexps {
				if v.Match([]byte(intf.Name)) {
					match = true
					break
				}
			}

			if !match {
				continue
			}
		}

		sn, exist := namesMap[g.mockName(intf.Name)]
		if exist {
			newMethods := make([]*model.Method, 0)
			for _, m := range intf.Methods {
				if _, exist = sn.Methods[m.Name]; exist {
					continue
				}
				newMethods = append(newMethods, m)
			}

			if len(newMethods) != 0 {
				intf.Methods = newMethods
				mockType := g.mockName(intf.Name)
				g.GenerateMockMethods(mockType, intf, outputPackagePath)
			}

			continue
		}

		if err := g.GenerateMockInterface(intf, outputPackagePath); err != nil {
			return err
		}
	}

	return nil
}

// The name of the mock type to use for the given interface identifier.
func (g *generator) mockName(typeName string) string {
	if mockName, ok := g.mockNames[typeName]; ok {
		return mockName
	}

	for _, suffix := range []string{"Interface", "Service", "Handler", "Server"} {
		if suffix == typeName[len(typeName)-len(suffix):] {
			typeName = typeName[:len(typeName)-len(suffix)]
			break
		}
	}

	return strings.ToLower(string(typeName[0])) + typeName[1:]
}

func (g *generator) GenerateMockInterface(intf *model.Interface, outputPackagePath string) error {
	mockType := g.mockName(intf.Name)

	g.p("")

	for _, doc := range intf.Doc {
		if strings.HasPrefix(strings.ToLower(doc), "//go:generate ") { // 生成语句不复制到实现文件中
			continue
		}

		g.p("%v", doc)
	}

	if 0 == len(intf.Comment) {
		g.p("type %v struct {", mockType)
	} else {
		g.p("type %v struct { // %v", mockType, intf.Comment)
	}
	g.in()
	g.out()
	g.p("}")
	g.p("")

	// TODO: Re-enable this if we can import the interface reliably.
	// g.p("// Verify that the mock satisfies the interface at compile time.")
	// g.p("var _ %v = (*%v)(nil)", typeName, mockType)
	// g.p("")

	g.p("// New%v create a new %v object", mockType, mockType)
	if 0 == len(intf.Comment) {
		g.p("func New%v(_ context.Context) *%v {", mockType, mockType)
	} else {
		g.p("func New%v(_ context.Context) *%v { // %v", mockType, mockType, intf.Comment)
	}

	g.in()
	g.p("obj := &%v{}", mockType)
	g.p("")
	g.p("// TODO: New%v(_ context.Context) Not implemented", mockType)
	g.p("")
	g.p("return obj")
	g.out()
	g.p("}")
	g.p("")

	g.GenerateMockMethods(mockType, intf, outputPackagePath)

	return nil
}

func (g *generator) GenerateMockMethods(mockType string, intf *model.Interface, pkgOverride string) {
	for _, m := range intf.Methods {
		g.p("")
		_ = g.GenerateMockMethod(mockType, m, pkgOverride)
	}
}

// GenerateMockMethod generates a mock method implementation.
// If non-empty, pkgOverride is the package in which unqualified types reside.
func (g *generator) GenerateMockMethod(mockType string, m *model.Method, pkgOverride string) error {
	argNames := g.getArgNames(m, pkgOverride)
	argTypes := g.getArgTypes(m, pkgOverride)
	argString := makeArgString(argNames, argTypes)

	rets := make([]string, len(m.Out))
	for i, p := range m.Out {
		rets[i] = p.Type.String(g.packageMap, pkgOverride)
	}
	retString := strings.Join(rets, ", ")
	if len(rets) > 1 {
		retString = "(" + retString + ")"
	}
	if retString != "" {
		retString = " " + retString
	}

	ia := newIdentifierAllocator(argNames)
	idRecv := ia.allocateIdentifier("m")

	for _, doc := range m.Doc {
		if strings.HasPrefix(strings.ToLower(doc), "//go:generate ") { // 生成语句不复制到实现文件中
			continue
		}

		g.p("%v", doc)
	}
	if 0 == len(m.Comment) {
		g.p("func (%v *%v) %v(%v)%v {", idRecv, mockType, m.Name, argString, retString)
	} else {
		g.pf("func (%v *%v) %v(%v)%v { // %v", idRecv, mockType, m.Name, argString, retString, m.Comment)
	}

	g.in()

	g.p("// TODO: %v.%v(%v)%v Not implemented", mockType, m.Name, argString, retString)
	g.p("")
	g.p("panic(\"%v.%v(%v)%v Not implemented\")", mockType, m.Name, argString, retString)
	g.out()
	g.p("}")
	return nil
}

func (g *generator) getArgNames(m *model.Method, pkgOverride string) []string {
	argNames := make([]string, len(m.In))
	args := make(map[string]bool)
	for i, p := range m.In {
		name := p.Name
		if name == "" || name == "_" {
			typName := p.Type.String(g.packageMap, pkgOverride)
			switch typName {
			case "error":
				name = "err"
			case "context.Context":
				name = "ctx"
			default:
				name = "arg"
			}

			if args[name] {
				name = name + strconv.Itoa(i)
			}
		}
		args[name] = true
		argNames[i] = name
	}
	if m.Variadic != nil {
		name := m.Variadic.Name
		if name == "" {
			name = fmt.Sprintf("arg%d", len(m.In))
		}
		argNames = append(argNames, name)
	}
	return argNames
}

func (g *generator) getArgTypes(m *model.Method, pkgOverride string) []string {
	argTypes := make([]string, len(m.In))
	for i, p := range m.In {
		argTypes[i] = p.Type.String(g.packageMap, pkgOverride)
	}
	if m.Variadic != nil {
		argTypes = append(argTypes, "..."+m.Variadic.Type.String(g.packageMap, pkgOverride))
	}
	return argTypes
}

// Output returns the generator's output, formatted in the standard Go style.
func (g *generator) Output() (n int, err error) {
	src, err := format.Source(g.buf.Bytes())
	if err != nil {
		log.Fatalf("Failed to format generated source code: %s\n%s", err, g.buf.String())
	}

	dst := os.Stdout
	if len(g.dstFileName) > 0 {
		if err := os.MkdirAll(filepath.Dir(*destination), os.ModePerm); err != nil {
			log.Fatalf("Unable to create directory: %v", err)
		}
		var f *os.File
		var err error
		if g.head {
			f, err = os.Create(*destination)
		} else {
			f, err = os.OpenFile(*destination, os.O_RDWR|os.O_APPEND, 0666)
		}

		if err != nil {
			log.Fatalf("Failed opening destination file: %v", err)
		}
		defer dst.Close()
		dst = f
	}

	return dst.Write(src)
}