From b56aba18fe5b4a509477955ce519813aa943cf75 Mon Sep 17 00:00:00 2001
From: sudo pacman -Syu <hauvipapro@gmail.com>
Date: Mon, 10 Jul 2023 20:41:08 +0700
Subject: [PATCH] feat: add stock mode (#37)

* feat: add stock flag

* feat: add formatter option stock

* feat: actually split stock import

* fix: missing flag when init cmd

* chore: add how to use stock mode

* chore: add stock mode in flag usage

* chore: remove dup

* chore: consistent docs
---
 README.md                            |  5 ++++
 internal/cli/action.go               |  3 ++
 internal/cli/app.go                  |  7 +++++
 internal/imports/formatter.go        | 43 ++++++++++++++++------------
 internal/imports/formatter_option.go |  6 ++++
 5 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 7b19e8a..e9603c2 100644
--- a/README.md
+++ b/README.md
@@ -29,6 +29,11 @@ go install github.com/haunt98/gofimports/cmd/gofimports@latest
 # - print diff (-d)
 # - company prefix, split using comma (,)
 gofimports -l -w -d --company github.com/make-go-great,github.com/haunt98 ./internal
+
+# Format ./internal with:
+# - write
+# - stock mode, only split standard and non standard
+gofimports -w --stock ./internal
 ```
 
 Example result:
diff --git a/internal/cli/action.go b/internal/cli/action.go
index 1204159..5e41a0a 100644
--- a/internal/cli/action.go
+++ b/internal/cli/action.go
@@ -19,6 +19,7 @@ type action struct {
 		diff          bool
 		verbose       bool
 		profiler      bool
+		stock         bool
 	}
 }
 
@@ -33,6 +34,7 @@ func (a *action) getFlags(c *cli.Context) {
 	a.flags.diff = c.Bool(flagDiffName)
 	a.flags.verbose = c.Bool(flagVerboseName)
 	a.flags.profiler = c.Bool(flagProfilerName)
+	a.flags.stock = c.Bool(flagStockName)
 
 	if a.flags.verbose {
 		fmt.Printf("flags: %+v\n", a.flags)
@@ -73,6 +75,7 @@ func (a *action) Run(c *cli.Context) error {
 		imports.FormatterWithDiff(a.flags.diff),
 		imports.FormatterWithVerbose(a.flags.verbose),
 		imports.FormatterWithCompanyPrefix(a.flags.companyPrefix),
+		imports.FormatterWithStock(a.flags.stock),
 	)
 	if err != nil {
 		return fmt.Errorf("imports: failed to new formatter: %w", err)
diff --git a/internal/cli/app.go b/internal/cli/app.go
index 36b4b58..425d19d 100644
--- a/internal/cli/app.go
+++ b/internal/cli/app.go
@@ -30,6 +30,9 @@ const (
 
 	flagProfilerName  = "profiler"
 	flagProfilerUsage = "go profiler, for debug only"
+
+	flagStockName  = "stock"
+	flagStockUsage = "stock mode, only split standard pkg and non standard, ignore company flag"
 )
 
 var (
@@ -77,6 +80,10 @@ func NewApp() *App {
 				Name:  flagProfilerName,
 				Usage: flagProfilerUsage,
 			},
+			&cli.BoolFlag{
+				Name:  flagStockName,
+				Usage: flagStockUsage,
+			},
 		},
 		Action: a.Run,
 	}
diff --git a/internal/imports/formatter.go b/internal/imports/formatter.go
index 5e20f1d..09e7ba1 100644
--- a/internal/imports/formatter.go
+++ b/internal/imports/formatter.go
@@ -62,6 +62,7 @@ type Formatter struct {
 	isWrite          bool
 	isDiff           bool
 	isVerbose        bool
+	isStock          bool
 }
 
 func NewFormmater(opts ...FormatterOptionFn) (*Formatter, error) {
@@ -324,10 +325,14 @@ func (ft *Formatter) groupDSTImportSpecs(importSpecs []*dst.ImportSpec, moduleNa
 	result := make(map[string][]*dst.ImportSpec)
 	result[stdImport] = make([]*dst.ImportSpec, 0, 8)
 	result[thirdPartyImport] = make([]*dst.ImportSpec, 0, 8)
-	if len(ft.companyPrefixes) != 0 {
-		result[companyImport] = make([]*dst.ImportSpec, 0, 8)
+	if !ft.isStock {
+		// Only split company, local imports if not stock
+		// Otherwise everything is third party imports
+		if len(ft.companyPrefixes) != 0 {
+			result[companyImport] = make([]*dst.ImportSpec, 0, 8)
+		}
+		result[localImport] = make([]*dst.ImportSpec, 0, 8)
 	}
-	result[localImport] = make([]*dst.ImportSpec, 0, 8)
 
 	for _, importSpec := range importSpecs {
 		// "github.com/abc/xyz" -> github.com/abc/xyz
@@ -338,23 +343,25 @@ func (ft *Formatter) groupDSTImportSpecs(importSpecs []*dst.ImportSpec, moduleNa
 			continue
 		}
 
-		if strings.HasPrefix(importPath, moduleName) {
-			result[localImport] = append(result[localImport], importSpec)
-			continue
-		}
-
-		if len(ft.companyPrefixes) != 0 {
-			existImport := false
-			for companyPrefix := range ft.companyPrefixes {
-				if strings.HasPrefix(importPath, companyPrefix) {
-					result[companyImport] = append(result[companyImport], importSpec)
-					existImport = true
-					break
-				}
+		if !ft.isStock {
+			if strings.HasPrefix(importPath, moduleName) {
+				result[localImport] = append(result[localImport], importSpec)
+				continue
 			}
 
-			if existImport {
-				continue
+			if len(ft.companyPrefixes) != 0 {
+				existImport := false
+				for companyPrefix := range ft.companyPrefixes {
+					if strings.HasPrefix(importPath, companyPrefix) {
+						result[companyImport] = append(result[companyImport], importSpec)
+						existImport = true
+						break
+					}
+				}
+
+				if existImport {
+					continue
+				}
 			}
 		}
 
diff --git a/internal/imports/formatter_option.go b/internal/imports/formatter_option.go
index 0e896f4..e9e9e7d 100644
--- a/internal/imports/formatter_option.go
+++ b/internal/imports/formatter_option.go
@@ -40,3 +40,9 @@ func FormatterWithCompanyPrefix(companyPrefix string) FormatterOptionFn {
 		}
 	}
 }
+
+func FormatterWithStock(isStock bool) FormatterOptionFn {
+	return func(ft *Formatter) {
+		ft.isStock = isStock
+	}
+}