improve performance of mock generation (#396)

golang · codyoss · Feb 22, 2020 · Feb 17, 2020 · Feb 22, 2020 · Feb 22, 2020
commit a53af468d5295e7fc1ff5b37d040ca945ee86fd0
diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
@@ -19,8 +19,8 @@ package main
 // TODO: This does not support embedding package-local interfaces in a separate file.
 
 import (
+	"bufio"
 	"bytes"
-	"encoding/json"
 	"flag"
 	"fmt"
 	"go/build"
@@ -304,10 +304,16 @@ func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPac
 	}
 	sort.Strings(sortedPaths)
 
+	importPaths := make([]string, 0)
+	for _, pth := range sortedPaths {
+		importPaths = append(importPaths, pth)
+	}
+	packagesName := lookupPackagesName(importPaths)
+
 	g.packageMap = make(map[string]string, len(im))
 	localNames := make(map[string]bool, len(im))
 	for _, pth := range sortedPaths {
-		base, ok := lookupPackageName(pth)
+		base, ok := packagesName[pth]
 		if !ok {
 			base = sanitize(path.Base(pth))
 		}
@@ -622,22 +628,23 @@ func (g *generator) Output() []byte {
 	return src
 }
 
-func lookupPackageName(importPath string) (string, bool) {
-	var pkg struct {
-		Name string
-	}
+func lookupPackagesName(importPaths []string) map[string]string {
+	names := make(map[string]string)
 	b := bytes.NewBuffer(nil)
-	cmd := exec.Command("go", "list", "-json", importPath)
+	args := []string{"list", "-f", "{{.Name}} {{.ImportPath}}"}
+	args = append(args, importPaths...)
+	cmd := exec.Command("go", args...)
 	cmd.Stdout = b
-	err := cmd.Run()
-	if err != nil {
-		return "", false
-	}
-	err = json.Unmarshal(b.Bytes(), &pkg)
-	if err != nil {
-		return "", false
+	cmd.Run()
+	sc := bufio.NewScanner(b)
+	for sc.Scan() {
+		pkg := strings.Split(sc.Text(), " ")
+		if len(pkg) != 2 {
+			log.Fatalln("Unable to lookup packages name: invalid output of go list command")
+		}
+		names[pkg[1]] = pkg[0]
 	}
-	return pkg.Name, true
+	return names
 }
 
 func printVersion() {

diff --git a/mockgen/mockgen_test.go b/mockgen/mockgen_test.go
@@ -335,30 +335,29 @@ func TestGetArgNames(t *testing.T) {
 	}
 }
 
-func Test_lookupPackageName(t *testing.T) {
-	type args struct {
-		importPath string
-	}
+func Test_lookupPackagesName(t *testing.T) {
 	tests := []struct {
-		name            string
 		importPath      string
 		wantPackageName string
-		wantOK          bool
+		shouldPresent   bool
 	}{
-		{"golang package", "context", "context", true},
-		{"third party", "golang.org/x/tools/present", "present", true},
-		{"modules", "rsc.io/quote/v3", "quote", true},
-		{"fail", "this/should/not/work", "", false},
+		{"context", "context", true},
+		{"golang.org/x/tools/present", "present", true},
+		{"rsc.io/quote/v3", "quote", true},
+		{"this/should/not/work", "", false},
+	}
+	importPaths := make([]string, 0)
+	for _, t := range tests {
+		importPaths = append(importPaths, t.importPath)
 	}
+	gotPackagesName := lookupPackagesName(importPaths)
 	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			gotPackageName, gotOk := lookupPackageName(tt.importPath)
-			if gotPackageName != tt.wantPackageName {
-				t.Errorf("lookupPackageName() gotPackageName = %v, wantPackageName %v", gotPackageName, tt.wantPackageName)
-			}
-			if gotOk != tt.wantOK {
-				t.Errorf("lookupPackageName() gotOk = %v, wantOK %v", gotOk, tt.wantOK)
-			}
-		})
+		gotPackageName, gotOk := gotPackagesName[tt.importPath]
+		if gotPackageName != tt.wantPackageName {
+			t.Errorf("lookupPackagesName() gotPackageName = %v, wantPackageName = %v", gotPackageName, tt.wantPackageName)
+		}
+		if gotOk != tt.shouldPresent {
+			t.Errorf("lookupPackageName() gotOk = %v, shouldPresent = %v", gotOk, tt.shouldPresent)
+		}
 	}
 }
diff --git a/mockgen/parse.go b/mockgen/parse.go
@@ -432,6 +432,15 @@ func (p *fileParser) parseType(pkg string, typ ast.Expr) (model.Type, error) {
 // importsOfFile returns a map of package name to import path
 // of the imports in file.
 func importsOfFile(file *ast.File) (normalImports map[string]string, dotImports []string) {
+	importPaths := make([]string, 0)
+	for _, is := range file.Imports {
+		if is.Name != nil {
+			continue
+		}
+		importPath := is.Path.Value[1 : len(is.Path.Value)-1] // remove quotes
+		importPaths = append(importPaths, importPath)
+	}
+	packagesName := lookupPackagesName(importPaths)
 	normalImports = make(map[string]string)
 	dotImports = make([]string, 0)
 	for _, is := range file.Imports {
@@ -445,7 +454,7 @@ func importsOfFile(file *ast.File) (normalImports map[string]string, dotImports
 			}
 			pkgName = is.Name.Name
 		} else {
-			pkg, ok := lookupPackageName(importPath)
+			pkg, ok := packagesName[importPath]
 			if !ok {
 				// Fallback to import path suffix. Note that this is uncertain.
 				_, last := path.Split(importPath)