Tests and fixes

README.md

@@ -58,7 +58,7 @@ package main
 import (
   "fmt"
   go "github.com/golang"
-  "github.com/daixiang0"
+  "github.com/daixiang0/gci"
 )
 ```
 
@@ -82,7 +82,7 @@ package main
 import (
   "fmt"
   _ "github.com/golang" // golang
-  "github.com/daixiang0"
+  "github.com/daixiang0/gci"
 )
 ```
 
@@ -108,7 +108,7 @@ import (
   "fmt"
   // golang
   _ "github.com/golang"
-  "github.com/daixiang0"
+  "github.com/daixiang0/gci"
 )
 ```
 
@@ -128,6 +128,5 @@ import (
 
 ## TODO
 
-- Support multi-3rd-party packages
 - Support multiple lines of comment in import block
-- Add testcases
+- Add more testcases

go.mod

@@ -2,4 +2,7 @@ module github.com/daixiang0/gci
 
 go 1.14
 
-require golang.org/x/tools v0.0.0-20201118003311-bd56c0adb394
+require (
+	github.com/stretchr/testify v1.7.0 // indirect
+	golang.org/x/tools v0.0.0-20201118003311-bd56c0adb394
+)

go.sum

@@ -1,3 +1,10 @@
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
@@ -22,3 +29,6 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/gci/gci.go

@@ -12,13 +12,17 @@ import (
 	"strings"
 )
 
+type PkgType int
+
 const (
 	// pkg type: standard, remote, local
-	standard int = iota
+	standard PkgType = iota
 	// 3rd-party packages
 	remote
 	local
+)
 
+const (
 	commentFlag = "//"
 )
 
@@ -37,7 +41,7 @@ type FlagSet struct {
 }
 
 type pkg struct {
-	list    map[int][]string
+	list    map[PkgType][]string
 	comment map[string]string
 	alias   map[string]string
 }
@@ -51,13 +55,10 @@ func ParseLocalFlag(str string) []string {
 }
 
 func newPkg(data [][]byte, localFlag []string) *pkg {
-	listMap := make(map[int][]string)
-	commentMap := make(map[string]string)
-	aliasMap := make(map[string]string)
 	p := &pkg{
-		list:    listMap,
-		comment: commentMap,
-		alias:   aliasMap,
+		list:    make(map[PkgType][]string),
+		comment: make(map[string]string),
+		alias:   make(map[string]string),
 	}
 
 	formatData := make([]string, 0)
@@ -111,51 +112,47 @@ func newPkg(data [][]byte, localFlag []string) *pkg {
 
 // fmt format import pkgs as expected
 func (p *pkg) fmt() []byte {
-	ret := make([]string, 0, 100)
+	var lines []string
 
-	for pkgType := range []int{standard, remote, local} {
+	for _, pkgType := range []PkgType{standard, remote, local} {
+		if len(p.list[pkgType]) == 0 {
+			continue
+		}
+		if len(lines) > 0 && lines[len(lines)-1] != "" {
+			lines = append(lines, "")
+		}
 		sort.Strings(p.list[pkgType])
 		for _, s := range p.list[pkgType] {
 			if p.comment[s] != "" {
-				l := fmt.Sprintf("%s%s%s%s", linebreak, indent, p.comment[s], linebreak)
-				ret = append(ret, l)
+				if len(lines) > 0 && lines[len(lines)-1] != "" {
+					lines = append(lines, "")
+				}
+				lines = append(lines, indent+p.comment[s])
 			}
 
 			if p.alias[s] != "" {
-				s = fmt.Sprintf("%s%s%s%s%s", indent, p.alias[s], blank, s, linebreak)
+				lines = append(lines, indent+p.alias[s]+blank+s)
 			} else {
-				s = fmt.Sprintf("%s%s%s", indent, s, linebreak)
+				lines = append(lines, indent+s)
 			}
-
-			ret = append(ret, s)
-		}
-
-		if len(p.list[pkgType]) > 0 {
-			ret = append(ret, linebreak)
 		}
 	}
-	if len(ret) > 0 && ret[len(ret)-1] == linebreak {
-		ret = ret[:len(ret)-1]
-	}
 
-	// remove duplicate empty lines
-	s1 := fmt.Sprintf("%s%s%s%s", linebreak, linebreak, linebreak, indent)
-	s2 := fmt.Sprintf("%s%s%s", linebreak, linebreak, indent)
-	return []byte(strings.ReplaceAll(strings.Join(ret, ""), s1, s2))
+	return []byte(strings.Join(lines, linebreak) + linebreak)
 }
 
 // getPkgInfo assume line is a import path, and return (path, alias, comment)
 func getPkgInfo(line string, comment bool) (string, string, string) {
 	if comment {
-		s := strings.Split(line, commentFlag)
-		pkgArray := strings.Split(s[0], blank)
+		s := strings.SplitN(line, commentFlag, 2)
+		pkgArray := strings.Fields(s[0])
 		if len(pkgArray) > 1 {
-			return pkgArray[1], pkgArray[0], fmt.Sprintf("%s%s%s", commentFlag, blank, strings.TrimSpace(s[1]))
+			return pkgArray[1], pkgArray[0], commentFlag + s[1]
 		} else {
-			return strings.TrimSpace(pkgArray[0]), "", fmt.Sprintf("%s%s%s", commentFlag, blank, strings.TrimSpace(s[1]))
+			return pkgArray[0], "", commentFlag + s[1]
 		}
 	} else {
-		pkgArray := strings.Split(line, blank)
+		pkgArray := strings.Fields(line)
 		if len(pkgArray) > 1 {
 			return pkgArray[1], pkgArray[0], ""
 		} else {
@@ -164,7 +161,7 @@ func getPkgInfo(line string, comment bool) (string, string, string) {
 	}
 }
 
-func getPkgType(line string, localFlag []string) int {
+func getPkgType(line string, localFlag []string) PkgType {
 	pkgName := strings.Trim(line, "\"\\`")
 
 	for _, localPkg := range localFlag {
@@ -363,7 +360,7 @@ func Run(filename string, set *FlagSet) ([]byte, []byte, error) {
 		return nil, nil, nil
 	}
 	end := bytes.Index(src[start:], importEndFlag) + start
-	
+
 	// in case import flags are part of a codegen template, or otherwise "wrong"
 	if start+len(importStartFlag) > end {
 		return nil, nil, nil

pkg/gci/gci_test.go

@@ -1,15 +1,21 @@
 package gci
 
 import (
+	"bytes"
 	"fmt"
+	"io/ioutil"
+	"path/filepath"
+	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestGetPkgType(t *testing.T) {
 	testCases := []struct {
 		Line           string
 		LocalFlag      string
-		ExpectedResult int
+		ExpectedResult PkgType
 	}{
 		{Line: `"foo/pkg/bar"`, LocalFlag: "", ExpectedResult: remote},
 		{Line: `"foo/pkg/bar"`, LocalFlag: "foo", ExpectedResult: local},
@@ -62,3 +68,49 @@ func TestGetPkgType(t *testing.T) {
 		})
 	}
 }
+
+func TestRun(t *testing.T) {
+	fileinfos, err := ioutil.ReadDir("testdata")
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, fileinfo := range fileinfos {
+		inname := fileinfo.Name()
+		if strings.HasPrefix(inname, ".") || !strings.HasSuffix(inname, ".in.go") {
+			continue
+		}
+		name := strings.TrimSuffix(inname, ".in.go")
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			expectedInput, err := ioutil.ReadFile(filepath.Join("testdata", inname))
+			if err != nil {
+				t.Fatal(err)
+			}
+			outname := name + ".out.go"
+			expectedOutput, err := ioutil.ReadFile(filepath.Join("testdata", outname))
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			flagSet := &FlagSet{
+				LocalFlag: []string{
+					"github.com/daixiang0",
+					"github.com/local",
+				},
+			}
+
+			actualInput, actualOutput, err := Run(filepath.Join("testdata", inname), flagSet)
+			assert.Equal(t, string(expectedInput), string(actualInput), "input")
+			if bytes.Equal(expectedInput, expectedOutput) {
+				assert.Nil(t, actualOutput, "output")
+
+			}
+			if actualOutput == nil {
+				actualOutput = actualInput
+			}
+			assert.Equal(t, string(expectedOutput), string(actualOutput), "output")
+			assert.NoError(t, err)
+		})
+	}
+}

pkg/gci/testdata/already-good.in.go

@@ -0,0 +1,8 @@
+package main
+import (
+	"fmt"
+
+	g "github.com/golang"
+
+	"github.com/daixiang0/gci"
+)

pkg/gci/testdata/already-good.out.go

@@ -0,0 +1,8 @@
+package main
+import (
+	"fmt"
+
+	g "github.com/golang"
+
+	"github.com/daixiang0/gci"
+)

pkg/gci/testdata/comment-with-slashslash.in.go

@@ -0,0 +1,5 @@
+package main
+
+import (
+	"fmt" // https://pkg.go.dev/fmt
+)

pkg/gci/testdata/comment-with-slashslash.out.go

@@ -0,0 +1,6 @@
+package main
+
+import (
+	// https://pkg.go.dev/fmt
+	"fmt"
+)

pkg/gci/testdata/leading-comment.in.go

@@ -0,0 +1,6 @@
+package main
+
+import (
+	// foo
+	"fmt"
+)

pkg/gci/testdata/leading-comment.out.go

@@ -0,0 +1,6 @@
+package main
+
+import (
+	// foo
+	"fmt"
+)

pkg/gci/testdata/nolint.in.go

@@ -0,0 +1,9 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/forbidden/pkg" //nolint:depguard
+
+	_ "github.com/daixiang0/gci" //nolint:depguard
+)

pkg/gci/testdata/nolint.out.go

@@ -0,0 +1,11 @@
+package main
+
+import (
+	"fmt"
+
+	//nolint:depguard
+	"github.com/forbidden/pkg"
+
+	//nolint:depguard
+	_ "github.com/daixiang0/gci"
+)

pkg/gci/testdata/simple-case.in.go

@@ -0,0 +1,8 @@
+package main
+import (
+  "golang.org/x/tools"
+
+  "fmt"
+
+  "github.com/daixiang0/gci"
+)

pkg/gci/testdata/simple-case.out.go

@@ -0,0 +1,8 @@
+package main
+import (
+	"fmt"
+
+	"golang.org/x/tools"
+
+	"github.com/daixiang0/gci"
+)

pkg/gci/testdata/whitespace-test.in.go

@@ -0,0 +1,7 @@
+package main
+
+import (
+  "fmt"
+  "github.com/golang" // golang
+  alias  "github.com/daixiang0/gci"
+)

pkg/gci/testdata/whitespace-test.out.go

@@ -0,0 +1,10 @@
+package main
+
+import (
+	"fmt"
+
+	// golang
+	"github.com/golang"
+
+	alias "github.com/daixiang0/gci"
+)

pkg/gci/testdata/with-above-comment-and-alias.in.go

@@ -0,0 +1,7 @@
+package main
+import (
+  "fmt"
+  // golang
+  _ "github.com/golang"
+  "github.com/daixiang0/gci"
+)