feature: 如果生成时文件已存在，则只生成新加的部分

Author: nieml

generator.go

@@ -0,0 +1,357 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"go/format"
+	"go/token"
+	"log"
+	"os"
+	"path"
+	"path/filepath"
+	"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
+	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)
+
+	if err != nil {
+		g.head = true
+		g.generateHead(pkg, outputPkgName, outputPackagePath)
+	} else {
+		namesMap := make(map[string]*model.Struct)
+		for _, sn := range dstPkg.StructNames {
+			namesMap[sn.Name] = sn
+		}
+
+		newInterfaces := make([]*model.Interface, 0)
+		for _, intf := range pkg.Interfaces {
+			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 0 != len(newMethods) {
+					intf.Methods = newMethods
+					mockType := g.mockName(intf.Name)
+					g.GenerateMockMethods(mockType, intf, outputPackagePath)
+				}
+			} else {
+				newInterfaces = append(newInterfaces, intf)
+			}
+		}
+
+		pkg.Interfaces = newInterfaces
+		fmt.Printf("%+v-%+v-%+v\n", dstPkg.Interfaces, namesMap, pkg.Interfaces)
+	}
+
+	return g.generate(pkg, outputPkgName, outputPackagePath)
+}
+
+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("")
+
+	// 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
+	}
+
+	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, outputPkgName string, outputPackagePath string) error {
+	for _, intf := range pkg.Interfaces {
+		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
+	}
+
+	suffix := "Interface"
+	if suffix == typeName[len(typeName)-len(suffix):] {
+		return typeName[:len(typeName)-len(suffix)]
+	}
+
+	return typeName
+}
+
+func (g *generator) GenerateMockInterface(intf *model.Interface, outputPackagePath string) error {
+	mockType := g.mockName(intf.Name)
+
+	g.p("")
+	for _, doc := range intf.Doc {
+		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("")
+
+	for _, doc := range intf.Doc {
+		g.p("%v", doc)
+	}
+	if 0 == len(intf.Comment) {
+		g.p("func New%v() *%v {", mockType, mockType)
+	} else {
+		g.p("func New%v() *%v { // %v", mockType, mockType, intf.Comment)
+	}
+
+	g.in()
+	g.p("interfaceImpl := &%v{}", mockType)
+	g.p("")
+	g.p("// TODO: ...")
+	g.p("")
+	g.p("return interfaceImpl")
+	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)
+	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 {
+		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("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) []string {
+	argNames := make([]string, len(m.In))
+	for i, p := range m.In {
+		name := p.Name
+		if name == "" || name == "_" {
+			name = fmt.Sprintf("arg%d", i)
+		}
+		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)
+}

go.mod

@@ -3,6 +3,7 @@ module github.com/ssoor/implgen
 go 1.13
 
 require (
+	github.com/gobuffalo/packr/v2 v2.8.0
 	github.com/golang/mock v1.4.3
 	golang.org/x/mod v0.3.0
 	golang.org/x/tools v0.0.0-20200612220849-54c614fe050c