imp/imp.go

package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/gobuffalo/here"
	"github.com/urfave/cli/v2"
	"go.jolheiser.com/beaver"
	"go.jolheiser.com/beaver/color"
	"go.jolheiser.com/regexp"
)

var (
	module   string
	importRe = regexp.MustCompile(`(?ms)import \(([^)]+)\)`)
	otherRe  = regexp.MustCompile(`(?:var|const|func)\s`)
)

func runImp(ctx *cli.Context) error {
	if ctx.Bool("verbose") {
		beaver.Console.Level = beaver.DEBUG
	}
	info, err := here.Current()
	if err != nil {
		return err
	}
	module = info.Module.Path
	beaver.Debugf("Current module: %s", module)

	var failed bool
	diffColor := color.New(color.FgGreen, color.BgHiBlack)
	if err := filepath.Walk(".", func(walkPath string, walkInfo os.FileInfo, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}

		walkPath = filepath.ToSlash(walkPath)

		if s := checkSkip(walkPath, walkInfo, ctx.StringSlice("exclude")); s != CHECK {
			if s == SKIPDIR {
				return filepath.SkipDir
			}
			return nil
		}

		beaver.Debugf("Checking file %s", walkPath)
		data, err := ioutil.ReadFile(walkPath)
		if err != nil {
			return err
		}

		importStart := importRe.FindIndex(data)
		if importStart == nil {
			return nil
		}

		otherStart := otherRe.FindIndex(data)
		if otherStart != nil && otherStart[0] < importStart[0] {
			return nil
		}

		groups := importRe.Groups(string(data))
		if groups.Index(0) == "" {
			return nil
		}

		imports := strings.Split(groups.Index(1), "\n")
		for idx, i := range imports {
			imports[idx] = strings.TrimSpace(i)
		}

		formatted := formatImportStmt(splitImports(imports))
		if ctx.Bool("write") {
			replaced := strings.Replace(string(data), groups.Index(0), formatted, 1)
			return ioutil.WriteFile(walkPath, []byte(replaced), walkInfo.Mode())
		}

		if !strings.EqualFold(groups.Index(0), formatted) {
			failed = true
			beaver.Infof("File: %s", diffColor.Format(walkPath))
			beaver.Infof("Expected:\n%s", diffColor.Format(formatted))
			beaver.Infof("Got:\n%s\n", diffColor.Format(groups.Index(0)))
		}
		return nil
	}); err != nil {
		return err
	}
	if failed {
		return errors.New("imports are formatted incorrectly; this can be fixed with the `--write` flag")
	}
	return nil
}

func checkSkip(walkPath string, walkInfo os.FileInfo, excludePaths []string) walkStatus {

	// Skip current directory
	if strings.EqualFold(walkPath, ".") {
		return SKIP
	}

	// Skip hidden paths (starting with ".")
	if strings.HasPrefix(walkPath, ".") {
		beaver.Debugf("Skipping hidden path %s", walkPath)
		if walkInfo.IsDir() {
			return SKIPDIR
		}
		return SKIP
	}

	// Skip directories
	if walkInfo.IsDir() {
		return SKIP
	}

	// Skip non-Go files
	if !strings.HasSuffix(walkInfo.Name(), ".go") {
		beaver.Debugf("Skipping non-Go file %s", walkPath)
		return SKIP
	}

	// Skip excluded directories
	for _, e := range excludePaths {
		if strings.HasPrefix(walkPath, e) {
			beaver.Debugf("Skipping file %s because of exclude path %s", walkPath, e)
			return SKIP
		}
	}
	return CHECK
}

func splitImports(imports []string) []importList {
	// 0 -> stdlib
	// 1 -> this module
	// 2 -> others
	split := []importList{{}, {}, {}}
	var comment string
	for _, imp := range imports {
		if imp == "" {
			continue
		}
		i := imp
		if strings.HasPrefix(i, `//`) {
			var nl string
			if comment != "" {
				nl = "\n"
			}
			comment = nl + i
			continue
		}

		var name string
		if !strings.HasPrefix(i, `"`) {
			parts := strings.Split(i, " ")
			name = parts[0]
			i = parts[1]
		}
		switch {
		case strings.HasPrefix(i, `"`+module):
			split[1] = append(split[1], importItem{comment, name, i})
		case strings.Contains(i, "."):
			split[2] = append(split[2], importItem{comment, name, i})
		default:
			split[0] = append(split[0], importItem{comment, name, i})
		}
		comment = ""
	}

	for _, s := range split {
		sort.Sort(s)
	}
	return split
}

func formatImportStmt(imports []importList) string {
	var decl string
	for _, imp := range imports {
		var pre string
		if decl != "" {
			pre = "\n\n\t"
		}
		if len(imp) > 0 {
			decl += pre + strings.Join(imp.StringSlice(), "\n\t")
		}
	}
	return fmt.Sprintf(`import (
	%s
)`, decl)
}