From 059bec968c61383b574810040ba9410712de36c5 Mon Sep 17 00:00:00 2001
From: Nishanth Shanmugham <nishanth.gerrard@gmail.com>
Date: Sat, 23 Dec 2017 03:05:04 -0500
Subject: [PATCH] cmd/goimports: support multiple comma-separated imports in
 -local flag

In cmd/goimports, allow for the -local flag to accept a comma-separated
list of import path prefixes. Also, update the imports package
accordingly to support this.

Fixes golang/go#19188

Change-Id: I083d584df8c3a77532f0f66e9c5d970960180e0d
Reviewed-on: https://go-review.googlesource.com/85397
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
---
 cmd/goimports/goimports.go |  2 +-
 imports/fix.go             | 18 +++++++--
 imports/fix_test.go        | 76 +++++++++++++++++++++++++++++---------
 3 files changed, 73 insertions(+), 23 deletions(-)

diff --git a/cmd/goimports/goimports.go b/cmd/goimports/goimports.go
index 16e30834f0..0ce85c9c32 100644
--- a/cmd/goimports/goimports.go
+++ b/cmd/goimports/goimports.go
@@ -47,7 +47,7 @@ var (
 
 func init() {
 	flag.BoolVar(&options.AllErrors, "e", false, "report all errors (not just the first 10 on different lines)")
-	flag.StringVar(&imports.LocalPrefix, "local", "", "put imports beginning with this string after 3rd-party packages")
+	flag.StringVar(&imports.LocalPrefix, "local", "", "put imports beginning with this string after 3rd-party packages; comma-separated list")
 }
 
 func report(err error) {
diff --git a/imports/fix.go b/imports/fix.go
index de6608790e..d792a78c75 100644
--- a/imports/fix.go
+++ b/imports/fix.go
@@ -32,16 +32,26 @@ var (
 	testMu  sync.RWMutex // guards globals reset by tests; used only if inTests
 )
 
-// LocalPrefix, if set, instructs Process to sort import paths with the given
-// prefix into another group after 3rd-party packages.
+// LocalPrefix is a comma-separated string of import path prefixes, which, if
+// set, instructs Process to sort the import paths with the given prefixes
+// into another group after 3rd-party packages.
 var LocalPrefix string
 
+func localPrefixes() []string {
+	if LocalPrefix != "" {
+		return strings.Split(LocalPrefix, ",")
+	}
+	return nil
+}
+
 // importToGroup is a list of functions which map from an import path to
 // a group number.
 var importToGroup = []func(importPath string) (num int, ok bool){
 	func(importPath string) (num int, ok bool) {
-		if LocalPrefix != "" && strings.HasPrefix(importPath, LocalPrefix) {
-			return 3, true
+		for _, p := range localPrefixes() {
+			if strings.HasPrefix(importPath, p) {
+				return 3, true
+			}
 		}
 		return
 	},
diff --git a/imports/fix_test.go b/imports/fix_test.go
index 83ceb02501..b8d51a3646 100644
--- a/imports/fix_test.go
+++ b/imports/fix_test.go
@@ -1420,19 +1420,21 @@ const Y = bar.X
 // Tests that the LocalPrefix option causes imports
 // to be added into a later group (num=3).
 func TestLocalPrefix(t *testing.T) {
-	defer func(s string) { LocalPrefix = s }(LocalPrefix)
-	LocalPrefix = "foo/"
-
-	testConfig{
-		gopathFiles: map[string]string{
-			"foo/bar/bar.go": "package bar \n const X = 1",
-		},
-	}.test(t, func(t *goimportTest) {
-		buf, err := Process(t.gopath+"/src/test/t.go", []byte("package main \n const Y = bar.X \n const _ = runtime.GOOS"), &Options{})
-		if err != nil {
-			t.Fatal(err)
-		}
-		const want = `package main
+	tests := []struct {
+		config      testConfig
+		localPrefix string
+		src         string
+		want        string
+	}{
+		{
+			config: testConfig{
+				gopathFiles: map[string]string{
+					"foo/bar/bar.go": "package bar \n const X = 1",
+				},
+			},
+			localPrefix: "foo/",
+			src:         "package main \n const Y = bar.X \n const _ = runtime.GOOS",
+			want: `package main
 
 import (
 	"runtime"
@@ -1442,11 +1444,49 @@ import (
 
 const Y = bar.X
 const _ = runtime.GOOS
-`
-		if string(buf) != want {
-			t.Errorf("Got:\n%s\nWant:\n%s", buf, want)
-		}
-	})
+`,
+		},
+		{
+			config: testConfig{
+				gopathFiles: map[string]string{
+					"example.org/pkg/pkg.go":          "package pkg \n const A = 1",
+					"foo/bar/bar.go":                  "package bar \n const B = 1",
+					"code.org/r/p/expproj/expproj.go": "package expproj \n const C = 1",
+				},
+			},
+			localPrefix: "example.org/pkg,foo/,code.org",
+			src:         "package main \n const X = pkg.A \n const Y = bar.B \n const Z = expproj.C \n const _ = runtime.GOOS",
+			want: `package main
+
+import (
+	"runtime"
+
+	"code.org/r/p/expproj"
+	"example.org/pkg"
+	"foo/bar"
+)
+
+const X = pkg.A
+const Y = bar.B
+const Z = expproj.C
+const _ = runtime.GOOS
+`,
+		},
+	}
+
+	for _, tt := range tests {
+		tt.config.test(t, func(t *goimportTest) {
+			defer func(s string) { LocalPrefix = s }(LocalPrefix)
+			LocalPrefix = tt.localPrefix
+			buf, err := Process(t.gopath+"/src/test/t.go", []byte(tt.src), &Options{})
+			if err != nil {
+				t.Fatal(err)
+			}
+			if string(buf) != tt.want {
+				t.Errorf("Got:\n%s\nWant:\n%s", buf, tt.want)
+			}
+		})
+	}
 }
 
 // Tests that running goimport on files in GOROOT (for people hacking