cmd/goimports, imports: add LocalPrefix flag & option

This allows the caller to indicate they want certain
import paths to sort into another group after 3rd-party
imports when added by goimports. For example, running
'goimports -local example.com/' might produce

    import (
        "database/sql"
        "io"
        "strconv"

        "golang.org/x/net/context"

        "example.com/foo/bar"
        "example.com/foo/baz"
    )

Updates golang/go#12420

Change-Id: If6d88599f6cca2f102313bce95ba6ac46ffec1fe

cmd/goimports/goimports.go

@@ -49,6 +49,7 @@ var (
 
 func init() {
 	flag.BoolVar(&options.AllErrors, "e", false, "report all errors (not just the first 10 on different lines)")
+	flag.StringVar(&options.LocalPrefix, "local", "", "put imports beginning with this string after 3rd-party packages")
 }
 
 func report(err error) {

imports/fix.go

@@ -34,24 +34,30 @@ var (
 
 // 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) {
+var importToGroup = []func(importPath string, opt *Options) (num int, ok bool){
+	func(importPath string, opt *Options) (num int, ok bool) {
+		if opt.LocalPrefix != "" && strings.HasPrefix(importPath, opt.LocalPrefix) {
+			return 3, true
+		}
+		return
+	},
+	func(importPath string, opt *Options) (num int, ok bool) {
 		if strings.HasPrefix(importPath, "appengine") {
 			return 2, true
 		}
 		return
 	},
-	func(importPath string) (num int, ok bool) {
+	func(importPath string, opt *Options) (num int, ok bool) {
 		if strings.Contains(importPath, ".") {
 			return 1, true
 		}
 		return
 	},
 }
 
-func importGroup(importPath string) int {
+func importGroup(importPath string, opt *Options) int {
 	for _, fn := range importToGroup {
-		if n, ok := fn(importPath); ok {
+		if n, ok := fn(importPath, opt); ok {
 			return n
 		}
 	}

imports/fix_test.go

@@ -1257,6 +1257,37 @@ 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) {
+	testConfig{
+		gopathFiles: map[string]string{
+			"foo/bar/bar.go": "package bar \n const X = 1",
+		},
+	}.test(t, func(t *goimportTest) {
+		opt := &Options{LocalPrefix: "foo/"}
+
+		buf, err := Process(t.gopath+"/src/test/t.go", []byte("package main \n const Y = bar.X \n const _ = runtime.GOOS"), opt)
+		if err != nil {
+			t.Fatal(err)
+		}
+		const want = `package main
+
+import (
+	"runtime"
+
+	"foo/bar"
+)
+
+const Y = bar.X
+const _ = runtime.GOOS
+`
+		if string(buf) != want {
+			t.Errorf("Got:\n%s\nWant:\n%s", buf, want)
+		}
+	})
+}
+
 // Tests that running goimport on files in GOROOT (for people hacking
 // on Go itself) don't cause the GOPATH to be scanned (which might be
 // much bigger).

imports/imports.go

@@ -35,6 +35,8 @@ type Options struct {
 	TabWidth  int  // Tab width (8 if nil *Options provided)
 
 	FormatOnly bool // Disable the insertion and deletion of imports
+
+	LocalPrefix string // Use later group (num=3) for import paths with this prefix
 }
 
 // Process formats and adjusts imports for the provided file.
@@ -61,7 +63,7 @@ func Process(filename string, src []byte, opt *Options) ([]byte, error) {
 		}
 	}
 
-	sortImports(fileSet, file)
+	sortImports(fileSet, file, opt)
 	imps := astutil.Imports(fileSet, file)
 
 	var spacesBefore []string // import paths we need spaces before
@@ -73,7 +75,7 @@ func Process(filename string, src []byte, opt *Options) ([]byte, error) {
 		lastGroup := -1
 		for _, importSpec := range impSection {
 			importPath, _ := strconv.Unquote(importSpec.Path.Value)
-			groupNum := importGroup(importPath)
+			groupNum := importGroup(importPath, opt)
 			if groupNum != lastGroup && lastGroup != -1 {
 				spacesBefore = append(spacesBefore, importPath)
 			}

imports/sortimports.go

@@ -15,7 +15,7 @@ import (
 
 // sortImports sorts runs of consecutive import lines in import blocks in f.
 // It also removes duplicate imports when it is possible to do so without data loss.
-func sortImports(fset *token.FileSet, f *ast.File) {
+func sortImports(fset *token.FileSet, f *ast.File, opt *Options) {
 	for i, d := range f.Decls {
 		d, ok := d.(*ast.GenDecl)
 		if !ok || d.Tok != token.IMPORT {
@@ -40,11 +40,11 @@ func sortImports(fset *token.FileSet, f *ast.File) {
 		for j, s := range d.Specs {
 			if j > i && fset.Position(s.Pos()).Line > 1+fset.Position(d.Specs[j-1].End()).Line {
 				// j begins a new run.  End this one.
-				specs = append(specs, sortSpecs(fset, f, d.Specs[i:j])...)
+				specs = append(specs, sortSpecs(fset, f, d.Specs[i:j], opt)...)
 				i = j
 			}
 		}
-		specs = append(specs, sortSpecs(fset, f, d.Specs[i:])...)
+		specs = append(specs, sortSpecs(fset, f, d.Specs[i:], opt)...)
 		d.Specs = specs
 
 		// Deduping can leave a blank line before the rparen; clean that up.
@@ -95,7 +95,7 @@ type posSpan struct {
 	End   token.Pos
 }
 
-func sortSpecs(fset *token.FileSet, f *ast.File, specs []ast.Spec) []ast.Spec {
+func sortSpecs(fset *token.FileSet, f *ast.File, specs []ast.Spec, opt *Options) []ast.Spec {
 	// Can't short-circuit here even if specs are already sorted,
 	// since they might yet need deduplication.
 	// A lone import, however, may be safely ignored.
@@ -144,7 +144,7 @@ func sortSpecs(fset *token.FileSet, f *ast.File, specs []ast.Spec) []ast.Spec {
 	// Reassign the import paths to have the same position sequence.
 	// Reassign each comment to abut the end of its spec.
 	// Sort the comments by new position.
-	sort.Sort(byImportSpec(specs))
+	sort.Sort(byImportSpec{specs, opt})
 
 	// Dedup. Thanks to our sorting, we can just consider
 	// adjacent pairs of imports.
@@ -179,30 +179,33 @@ func sortSpecs(fset *token.FileSet, f *ast.File, specs []ast.Spec) []ast.Spec {
 	return specs
 }
 
-type byImportSpec []ast.Spec // slice of *ast.ImportSpec
+type byImportSpec struct {
+	a   []ast.Spec // slice of *ast.ImportSpec
+	opt *Options
+}
 
-func (x byImportSpec) Len() int      { return len(x) }
-func (x byImportSpec) Swap(i, j int) { x[i], x[j] = x[j], x[i] }
+func (x byImportSpec) Len() int      { return len(x.a) }
+func (x byImportSpec) Swap(i, j int) { x.a[i], x.a[j] = x.a[j], x.a[i] }
 func (x byImportSpec) Less(i, j int) bool {
-	ipath := importPath(x[i])
-	jpath := importPath(x[j])
+	ipath := importPath(x.a[i])
+	jpath := importPath(x.a[j])
 
-	igroup := importGroup(ipath)
-	jgroup := importGroup(jpath)
+	igroup := importGroup(ipath, x.opt)
+	jgroup := importGroup(jpath, x.opt)
 	if igroup != jgroup {
 		return igroup < jgroup
 	}
 
 	if ipath != jpath {
 		return ipath < jpath
 	}
-	iname := importName(x[i])
-	jname := importName(x[j])
+	iname := importName(x.a[i])
+	jname := importName(x.a[j])
 
 	if iname != jname {
 		return iname < jname
 	}
-	return importComment(x[i]) < importComment(x[j])
+	return importComment(x.a[i]) < importComment(x.a[j])
 }
 
 type byCommentPos []*ast.CommentGroup