Improve performance of mock generation (#397)

plakhotnii · web-flow · commit b48cb6623c04 · 2020-02-22T14:34:44.000-07:00
Optimizes resolving correct package name with `go list` command. Instead of calling go list for each import in file, now we do it for all imports in one call.
diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
@@ -304,10 +304,12 @@ func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPac
 	}
 	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 := lookupPackageName(pth)
+		base, ok := packagesName[pth]
 		if !ok {
 			base = sanitize(path.Base(pth))
 		}
@@ -622,22 +624,30 @@ func (g *generator) Output() []byte {
 	return src
 }
 
-func lookupPackageName(importPath string) (string, bool) {
+// createPackageMap returns a map of import path to package name
+// for specified importPaths.
+func createPackageMap(importPaths []string) map[string]string {
 	var pkg struct {
-		Name string
+		Name       string
+		ImportPath string
 	}
+	pkgMap := make(map[string]string)
 	b := bytes.NewBuffer(nil)
-	cmd := exec.Command("go", "list", "-json", importPath)
+	args := []string{"list", "-json"}
+	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()
+	dec := json.NewDecoder(b)
+	for dec.More() {
+		err := dec.Decode(&pkg)
+		if err != nil {
+			log.Printf("failed to decode 'go list' output: %v", err)
+			continue
+		}
+		pkgMap[pkg.ImportPath] = pkg.Name
 	}
-	return pkg.Name, true
+	return pkgMap
 }
 
 func printVersion() {
diff --git a/mockgen/mockgen_test.go b/mockgen/mockgen_test.go
@@ -335,10 +335,7 @@ func TestGetArgNames(t *testing.T) {
 	}
 }
 
-func Test_lookupPackageName(t *testing.T) {
-	type args struct {
-		importPath string
-	}
+func Test_createPackageMap(t *testing.T) {
 	tests := []struct {
 		name            string
 		importPath      string
@@ -350,14 +347,19 @@ func Test_lookupPackageName(t *testing.T) {
 		{"modules", "rsc.io/quote/v3", "quote", true},
 		{"fail", "this/should/not/work", "", false},
 	}
+	var importPaths []string
+	for _, t := range tests {
+		importPaths = append(importPaths, t.importPath)
+	}
+	packages := createPackageMap(importPaths)
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			gotPackageName, gotOk := lookupPackageName(tt.importPath)
+			gotPackageName, gotOk := packages[tt.importPath]
 			if gotPackageName != tt.wantPackageName {
-				t.Errorf("lookupPackageName() gotPackageName = %v, wantPackageName %v", gotPackageName, tt.wantPackageName)
+				t.Errorf("createPackageMap() gotPackageName = %v, wantPackageName = %v", gotPackageName, tt.wantPackageName)
 			}
 			if gotOk != tt.wantOK {
-				t.Errorf("lookupPackageName() gotOk = %v, wantOK %v", gotOk, tt.wantOK)
+				t.Errorf("createPackageMap() gotOk = %v, wantOK = %v", gotOk, tt.wantOK)
 			}
 		})
 	}
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) {
+	var importPaths []string
+	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 := createPackageMap(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)

Original file line number	Diff line number	Diff line change
`@@ -335,10 +335,7 @@ func TestGetArgNames(t *testing.T) {`
`335`	`335`	`}`
`336`	`336`	`}`
`337`	`337`
`338`		`-func Test_lookupPackageName(t *testing.T) {`
`339`		`- type args struct {`
`340`		`- importPath string`
`341`		`- }`
	`338`	`+func Test_createPackageMap(t *testing.T) {`
`342`	`339`	`tests := []struct {`
`343`	`340`	`name string`
`344`	`341`	`importPath string`
`@@ -350,14 +347,19 @@ func Test_lookupPackageName(t *testing.T) {`
`350`	`347`	`{"modules", "rsc.io/quote/v3", "quote", true},`
`351`	`348`	`{"fail", "this/should/not/work", "", false},`
`352`	`349`	`}`
	`350`	`+ var importPaths []string`
	`351`	`+ for _, t := range tests {`
	`352`	`+ importPaths = append(importPaths, t.importPath)`
	`353`	`+ }`
	`354`	`+ packages := createPackageMap(importPaths)`
`353`	`355`	`for _, tt := range tests {`
`354`	`356`	`t.Run(tt.name, func(t *testing.T) {`
`355`		`- gotPackageName, gotOk := lookupPackageName(tt.importPath)`
	`357`	`+ gotPackageName, gotOk := packages[tt.importPath]`
`356`	`358`	`if gotPackageName != tt.wantPackageName {`
`357`		`- t.Errorf("lookupPackageName() gotPackageName = %v, wantPackageName %v", gotPackageName, tt.wantPackageName)`
	`359`	`+ t.Errorf("createPackageMap() gotPackageName = %v, wantPackageName = %v", gotPackageName, tt.wantPackageName)`
`358`	`360`	`}`
`359`	`361`	`if gotOk != tt.wantOK {`
`360`		`- t.Errorf("lookupPackageName() gotOk = %v, wantOK %v", gotOk, tt.wantOK)`
	`362`	`+ t.Errorf("createPackageMap() gotOk = %v, wantOK = %v", gotOk, tt.wantOK)`
`361`	`363`	`}`
`362`	`364`	`})`
`363`	`365`	`}`