gofimports/internal/imports/formatter.go

488 lines
11 KiB
Go

package imports
import (
"bytes"
"errors"
"fmt"
"go/parser"
"go/token"
"io/fs"
"os"
"path/filepath"
"strings"
"sync"
"github.com/dave/dst"
"github.com/dave/dst/decorator"
"github.com/pkg/diff"
"golang.org/x/mod/modfile"
"golang.org/x/sync/errgroup"
"golang.org/x/tools/go/packages"
)
const (
// Use for group imports
stdImport = "std"
thirdPartyImport = "third-party"
companyImport = "company"
localImport = "local"
)
var (
ErrEmptyPaths = errors.New("empty paths")
ErrNotGoFile = errors.New("not go file")
ErrGoGeneratedFile = errors.New("go generated file")
ErrAlreadyFormatted = errors.New("already formatted")
ErrEmptyImport = errors.New("empty import")
ErrGoModNotExist = errors.New("go mod not exist")
ErrGoModEmptyModule = errors.New("go mod empty module")
ErrNotBytesBuffer = errors.New("not bytes.Buffer")
ErrNotDSTGenDecl = errors.New("not dst.GenDecl")
)
// https://pkg.go.dev/sync#Pool
var bufPool = sync.Pool{
New: func() any {
return new(bytes.Buffer)
},
}
// stdPackages -> save std packages for later search.
// moduleNames -> map path to its go.mod module name.
// formattedPaths -> make sure we not format path more than 1 time.
type Formatter struct {
stdPackages map[string]struct{}
moduleNames map[string]string
formattedPaths map[string]struct{}
companyPrefixes map[string]struct{}
eg errgroup.Group
muModuleNames sync.RWMutex
muFormattedPaths sync.RWMutex
isList bool
isWrite bool
isDiff bool
isVerbose bool
isStock bool
}
func NewFormmater(opts ...FormatterOptionFn) (*Formatter, error) {
ft := &Formatter{}
for _, opt := range opts {
opt(ft)
}
stdPackages, err := packages.Load(nil, "std")
if err != nil {
return nil, fmt.Errorf("packages: failed to load std: %w", err)
}
ft.stdPackages = make(map[string]struct{})
for _, stdPackage := range stdPackages {
ft.stdPackages[stdPackage.PkgPath] = struct{}{}
}
ft.moduleNames = make(map[string]string)
ft.formattedPaths = make(map[string]struct{})
return ft, nil
}
// Accept a list of files or directories
func (ft *Formatter) Format(paths ...string) error {
if len(paths) == 0 {
return ErrEmptyPaths
}
// Logic switch case copy from goimports, gofumpt
for _, path := range paths {
path := strings.TrimSpace(path)
if path == "" {
continue
}
switch dir, err := os.Stat(path); {
case err != nil:
return fmt.Errorf("os: failed to stat: [%s] %w", path, err)
case dir.IsDir():
if err := ft.formatDir(path); err != nil {
return err
}
default:
ft.eg.Go(func() error {
if err := ft.formatFile(path); err != nil {
if ft.isIgnoreError(err) {
return nil
}
return err
}
return nil
})
}
}
if err := ft.eg.Wait(); err != nil {
return err
}
return nil
}
// Copy from gofumpt
func (ft *Formatter) formatDir(path string) error {
if err := filepath.WalkDir(path, func(path string, dirEntry fs.DirEntry, err error) error {
if filepath.Base(path) == "vendor" {
return filepath.SkipDir
}
if err != nil {
return err
}
if dirEntry.IsDir() {
// Get module name ASAP to cache it
moduleName, err := ft.moduleName(path)
if err != nil {
return err
}
ft.log("formatFile: moduleName: [%s]\n", moduleName)
return nil
}
ft.eg.Go(func() error {
if err := ft.formatFile(path); err != nil {
if ft.isIgnoreError(err) {
return nil
}
return err
}
return nil
})
return nil
}); err != nil {
return fmt.Errorf("filepath: failed to walk dir: [%s] %w", path, err)
}
return nil
}
func (ft *Formatter) formatFile(path string) error {
ft.muFormattedPaths.RLock()
if _, ok := ft.formattedPaths[path]; ok {
ft.muFormattedPaths.RUnlock()
return nil
}
ft.muFormattedPaths.RUnlock()
// Return if not go file
if !isGoFile(filepath.Base(path)) {
return ErrNotGoFile
}
// Read file first
pathBytes, err := os.ReadFile(path)
if err != nil {
return fmt.Errorf("os: failed to read file: [%s] %w", path, err)
}
// Get module name of path
moduleName, err := ft.moduleName(path)
if err != nil {
return err
}
ft.log("formatFile: moduleName: [%s]\n", moduleName)
formattedBytes, err := ft.formatImports(path, pathBytes, moduleName)
if err != nil {
return err
}
if bytes.Equal(pathBytes, formattedBytes) {
return ErrAlreadyFormatted
}
if ft.isList {
fmt.Println("Formatted: ", path)
}
if ft.isWrite {
if err := os.WriteFile(path, formattedBytes, 0o600); err != nil {
return fmt.Errorf("os: failed to write file: [%s] %w", path, err)
}
}
if ft.isDiff {
if err := diff.Text(path+" before", path+" after", pathBytes, formattedBytes, os.Stdout); err != nil {
return fmt.Errorf("diff: failed to slices: %w", err)
}
}
ft.muFormattedPaths.Lock()
ft.formattedPaths[path] = struct{}{}
ft.muFormattedPaths.Unlock()
return nil
}
// Copy from goimports, gofumpt, goimports-reviser.
// First parse ast.
// Then group imports.
// Then format imports.
// Then print to bytes.
func (ft *Formatter) formatImports(
path string,
pathBytes []byte,
moduleName string,
) ([]byte, error) {
// Parse ast
fset := token.NewFileSet()
// Copy from gofumpt
parserMode := parser.Mode(0)
parserMode |= parser.ParseComments
parserMode |= parser.SkipObjectResolution
astFile, err := parser.ParseFile(fset, path, pathBytes, parserMode)
if err != nil {
return nil, fmt.Errorf("parser: failed to parse file [%s]: %w", path, err)
}
// Ignore generated file
if isGoGenerated(astFile) {
return nil, ErrGoGeneratedFile
}
dec := decorator.NewDecorator(fset)
dstFile, err := dec.DecorateFile(astFile)
if err != nil {
return nil, fmt.Errorf("decorator: failed to parse file [%s]: %w", path, err)
}
if len(dstFile.Imports) == 0 || len(dstFile.Decls) == 0 {
return nil, ErrEmptyImport
}
ft.logDSTImportSpecs("formatImports: dstImportSpecs", dstFile.Imports)
groupedDSTImportSpecs, err := ft.groupDSTImportSpecs(dstFile.Imports, moduleName)
if err != nil {
return nil, err
}
formattedDSTImportSpecs, err := ft.formatDSTImportSpecs(groupedDSTImportSpecs)
if err != nil {
return nil, err
}
ft.logDSTImportSpecs("formatImports: formattedDSTImportSpecs: ", formattedDSTImportSpecs)
// First update
dstFile.Imports = formattedDSTImportSpecs
genDecl, ok := dstFile.Decls[0].(*dst.GenDecl)
if !ok {
return nil, ErrNotDSTGenDecl
}
formattedGenSpecs := make([]dst.Spec, 0, len(genDecl.Specs))
// Append all imports first
for _, importSpec := range formattedDSTImportSpecs {
formattedGenSpecs = append(formattedGenSpecs, importSpec)
}
// Append all non imports later
for _, genSpec := range genDecl.Specs {
if _, ok := genSpec.(*dst.ImportSpec); !ok {
formattedGenSpecs = append(formattedGenSpecs, genSpec)
continue
}
}
// Second update
genDecl.Specs = formattedGenSpecs
b, ok := bufPool.Get().(*bytes.Buffer)
if !ok {
return nil, ErrNotBytesBuffer
}
b.Reset()
defer bufPool.Put(b)
if err := decorator.Fprint(b, dstFile); err != nil {
return nil, fmt.Errorf("decorator: failed to fprint [%s]: %w", path, err)
}
result := make([]byte, b.Len())
copy(result, b.Bytes())
return result, nil
}
func (ft *Formatter) groupDSTImportSpecs(importSpecs []*dst.ImportSpec, moduleName string) (map[string][]*dst.ImportSpec, error) {
result := make(map[string][]*dst.ImportSpec)
result[stdImport] = make([]*dst.ImportSpec, 0, 8)
result[thirdPartyImport] = 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)
}
for _, importSpec := range importSpecs {
// "github.com/abc/xyz" -> github.com/abc/xyz
importPath := strings.Trim(importSpec.Path.Value, `"`)
if _, ok := ft.stdPackages[importPath]; ok {
result[stdImport] = append(result[stdImport], importSpec)
continue
}
if !ft.isStock {
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 existImport {
continue
}
}
}
result[thirdPartyImport] = append(result[thirdPartyImport], importSpec)
}
ft.logDSTImportSpecs("groupDSTImportSpecs: stdImport", result[stdImport])
ft.logDSTImportSpecs("groupDSTImportSpecs: thirdPartyImport", result[thirdPartyImport])
if len(ft.companyPrefixes) != 0 {
ft.logDSTImportSpecs("groupDSTImportSpecs: companyImport", result[companyImport])
}
ft.logDSTImportSpecs("groupDSTImportSpecs: localImport", result[localImport])
return result, nil
}
func (ft *Formatter) formatDSTImportSpecs(groupedImportSpecs map[string][]*dst.ImportSpec,
) ([]*dst.ImportSpec, error) {
result := make([]*dst.ImportSpec, 0, 32)
appendToResultFn := func(groupImportType string) {
importSpecs, ok := groupedImportSpecs[groupImportType]
if !ok || len(importSpecs) == 0 {
return
}
for _, importSpec := range importSpecs {
importSpec.Decs.Before = dst.NewLine
importSpec.Decs.After = dst.NewLine
}
importSpecs[len(importSpecs)-1].Decs.After = dst.EmptyLine
result = append(result, importSpecs...)
}
appendToResultFn(stdImport)
appendToResultFn(thirdPartyImport)
appendToResultFn(companyImport)
appendToResultFn(localImport)
if len(result) == 0 {
return result, nil
}
result[len(result)-1].Decs.After = dst.NewLine
return result, nil
}
// Copy from goimports-reviser
// Get module name from go.mod of path
// If current path doesn't have go.mod, recursive find its parent path
func (ft *Formatter) moduleName(path string) (string, error) {
ft.muModuleNames.RLock()
if pkgName, ok := ft.moduleNames[path]; ok {
ft.muModuleNames.RUnlock()
return pkgName, nil
}
ft.muModuleNames.RUnlock()
// Copy from goimports-reviser
// Check path/go.mod first
// If not exist -> check ../go.mod
// Assume path is dir path, maybe wrong but it is ok for now
dirPath := filepath.Clean(path)
var goModPath string
for {
ft.muModuleNames.RLock()
if pkgName, ok := ft.moduleNames[dirPath]; ok {
ft.muModuleNames.RUnlock()
return pkgName, nil
}
ft.muModuleNames.RUnlock()
goModPath = filepath.Join(dirPath, "go.mod")
fileInfo, err := os.Stat(goModPath)
if err == nil && !fileInfo.IsDir() {
break
}
// Check ..
if dirPath == filepath.Dir(dirPath) {
// Reach root
break
}
dirPath = filepath.Dir(dirPath)
}
if goModPath == "" {
return "", ErrGoModNotExist
}
ft.log("moduleName: goModPath: %+v\n", goModPath)
goModPathBytes, err := os.ReadFile(goModPath)
if err != nil {
return "", fmt.Errorf("os: failed to read file: [%s] %w", goModPath, err)
}
goModFile, err := modfile.Parse(goModPath, goModPathBytes, nil)
if err != nil {
return "", fmt.Errorf("modfile: failed to parse: [%s] %w", goModPath, err)
}
result := goModFile.Module.Mod.Path
if result == "" {
return "", ErrGoModEmptyModule
}
ft.muModuleNames.Lock()
ft.moduleNames[path] = result
ft.muModuleNames.Unlock()
return result, nil
}
func (ft *Formatter) isIgnoreError(err error) bool {
return errors.Is(err, ErrNotGoFile) ||
errors.Is(err, ErrGoGeneratedFile) ||
errors.Is(err, ErrAlreadyFormatted) ||
errors.Is(err, ErrEmptyImport)
}