Auto-format whitespace between imports

3 files changed, 92 insertions, 11 deletions

diff --git a/internal/cmd/bracefmt/brace.go b/internal/cmd/gitalyfmt/format.go
index e6f55d1d4..1c3b2377d 100644
--- a/internal/cmd/bracefmt/brace.go
+++ b/internal/cmd/gitalyfmt/format.go
@@ -2,8 +2,10 @@ package main
 
 import (
 	"bytes"
+	"errors"
 	"go/scanner"
 	"go/token"
+	"regexp"
 )
 
 type editCommand int
@@ -19,20 +21,26 @@ type edit struct {
 	cmd  editCommand
 }
 
-func braceFmt(src []byte) []byte {
+var (
+	nonStdlibImportRegex = regexp.MustCompile(`^[^/]+\..*/`)
+)
+
+func format(src []byte) ([]byte, error) {
 	var s scanner.Scanner
 	fset := token.NewFileSet()
 	file := fset.AddFile("", fset.Base(), len(src))
 	s.Init(file, src, nil, scanner.ScanComments)
 
 	var (
-		edits                            []edit
-		lastNonEmptyLine, lastLBraceLine int
-		lastOuterRBraceLine              = -1
+		edits                               []edit
+		lastNonEmptyLine, lastLBraceLine    int
+		lastOuterRBraceLine                 = -1
+		importLine, lastNonStdlibImportLine int
+		inImports                           bool
 	)
 
 	for {
-		pos, tok, _ := s.Scan()
+		pos, tok, lit := s.Scan()
 		if tok == token.EOF {
 			break
 		}
@@ -42,6 +50,26 @@ func braceFmt(src []byte) []byte {
 		var nextEdit edit
 
 		switch tok {
+		case token.IMPORT:
+			importLine = currentLine
+		case token.LPAREN:
+			if currentLine == importLine {
+				inImports = true
+			}
+		case token.RPAREN:
+			inImports = false
+		case token.STRING:
+			if inImports {
+				if nonStdlibImportRegex.MatchString(lit) {
+					if currentLine-lastNonStdlibImportLine > 1 && lastNonStdlibImportLine > 0 {
+						nextEdit = edit{line: currentLine - 1, cmd: removeLine}
+					}
+					lastNonStdlibImportLine = currentLine
+				} else if lastNonStdlibImportLine > 0 {
+					// It would be nicer to fix this automatically, but how?
+					return nil, errors.New("stdlib import after non-stdlib import is not allowed")
+				}
+			}
 		case token.RBRACE:
 			if currentLine-lastNonEmptyLine > 1 {
 				// ......foo
@@ -101,5 +129,5 @@ func braceFmt(src []byte) []byte {
 		}
 	}
 
-	return bytes.Join(srcLines, []byte("\n"))
+	return bytes.Join(srcLines, []byte("\n")), nil
 }
diff --git a/internal/cmd/bracefmt/brace_test.go b/internal/cmd/gitalyfmt/format_test.go
index b3de996dd..4dae5ea83 100644
--- a/internal/cmd/bracefmt/brace_test.go
+++ b/internal/cmd/gitalyfmt/format_test.go
@@ -6,12 +6,13 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestBraceFmt(t *testing.T) {
+func TestFormat(t *testing.T) {
 	testCases := []struct {
 		desc      string
 		in        string
 		out       string
 		unchanged bool
+		fail      bool
 	}{
 		{
 			desc: "empty lines inside braces",
@@ -114,11 +115,60 @@ func foo() {
 `,
 			unchanged: true,
 		},
+		{
+			desc: "empty lines between non-stdlib imports",
+			in: `package main
+
+import (
+	"net/http"
+
+	"example.com/foo"
+
+	bar "example.com/bar"
+)
+
+func main() {}
+`,
+			out: `package main
+
+import (
+	"net/http"
+
+	"example.com/foo"
+	bar "example.com/bar"
+)
+
+func main() {}
+`,
+		},
+		{
+			desc: "alternating stdlib and non stdlib",
+			in: `package main
+
+import (
+	"net/http"
+
+	"example.com/foo"
+
+	"io"
+)
+
+func main() {}
+`,
+			fail: true,
+		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			out := braceFmt([]byte(tc.in))
+			out, err := format([]byte(tc.in))
+
+			if tc.fail {
+				require.Error(t, err, "expect format error")
+				return
+			}
+
+			require.NoError(t, err, "format error")
 
 			if tc.unchanged {
 				require.Equal(t, tc.in, string(out))
diff --git a/internal/cmd/bracefmt/main.go b/internal/cmd/gitalyfmt/main.go
index ffa4b05d9..bbc6b2d7f 100644
--- a/internal/cmd/bracefmt/main.go
+++ b/internal/cmd/gitalyfmt/main.go
@@ -9,7 +9,7 @@ import (
 )
 
 const (
-	progName = "bracefmt"
+	progName = "gitalyfmt"
 )
 
 var (
@@ -25,7 +25,7 @@ func main() {
 	}
 
 	if err := _main(flag.Args()); err != nil {
-		fmt.Fprintf(os.Stderr, "%s: fatal: %v", progName, err)
+		fmt.Fprintf(os.Stderr, "%s: fatal: %v\n", progName, err)
 		os.Exit(1)
 	}
 }
@@ -42,7 +42,10 @@ func _main(args []string) error {
 			return err
 		}
 
-		dst := braceFmt(src)
+		dst, err := format(src)
+		if err != nil {
+			return fmt.Errorf("%s: %v", f, err)
+		}
 		if bytes.Equal(src, dst) {
 			continue
 		}