mukan-ignite/ignite/pkg/xast/import.go

package xast

import (
	"bytes"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"strconv"

	"golang.org/x/tools/go/ast/astutil"

	"github.com/ignite/cli/v29/ignite/pkg/errors"
)

type (
	// importOpts represent the options for imp.
	importOpts struct {
		imports []imp
	}

	// ImportOptions configures code generation.
	ImportOptions func(*importOpts)

	imp struct {
		path string
		name string
	}
)

// WithImport add a new import at the end of the imports.
func WithImport(repo string) ImportOptions {
	return func(c *importOpts) {
		c.imports = append(c.imports, imp{
			path: repo,
			name: "",
		})
	}
}

// WithNamedImport add a new import with name at the end of the imports.
func WithNamedImport(name, repo string) ImportOptions {
	return func(c *importOpts) {
		c.imports = append(c.imports, imp{
			name: name,
			path: repo,
		})
	}
}

func newImportOptions() importOpts {
	return importOpts{
		imports: make([]imp, 0),
	}
}

// AppendImports appends import statements to the existing import block in Go source code content.
func AppendImports(fileContent string, imports ...ImportOptions) (string, error) {
	// apply global options.
	opts := newImportOptions()
	for _, o := range imports {
		o(&opts)
	}
	if len(opts.imports) == 0 {
		return fileContent, nil
	}

	fileSet := token.NewFileSet()

	// Parse the Go source code content.
	f, err := parser.ParseFile(fileSet, "", fileContent, parser.ParseComments)
	if err != nil {
		return "", err
	}
	cmap := ast.NewCommentMap(fileSet, f, f.Comments)

	// Add new import statements.
	for _, importPath := range opts.imports {
		deleteImportsByPath(fileSet, f, importPath.path)

		if !astutil.AddNamedImport(fileSet, f, importPath.name, importPath.path) {
			if hasImport(f, importPath.name, importPath.path) {
				continue
			}
			return "", errors.Errorf("failed to add import %s - %s", importPath.name, importPath.path)
		}
	}
	ast.SortImports(fileSet, f)

	f.Comments = cmap.Filter(f).Comments()

	// Format the modified AST.
	var buf bytes.Buffer
	if err := format.Node(&buf, fileSet, f); err != nil {
		return "", err
	}

	return buf.String(), nil
}

// RemoveImports removes import statements from the existing import block in Go source code content.
func RemoveImports(fileContent string, imports ...ImportOptions) (string, error) {
	// apply global options.
	opts := newImportOptions()
	for _, o := range imports {
		o(&opts)
	}
	if len(opts.imports) == 0 {
		return fileContent, nil
	}

	fileSet := token.NewFileSet()

	// Parse the Go source code content.
	f, err := parser.ParseFile(fileSet, "", fileContent, parser.ParseComments)
	if err != nil {
		return "", err
	}
	cmap := ast.NewCommentMap(fileSet, f, f.Comments)

	// Remove import statements.
	for _, importPath := range opts.imports {
		astutil.DeleteNamedImport(fileSet, f, importPath.name, importPath.path)
	}

	f.Comments = cmap.Filter(f).Comments()

	// Format the modified AST.
	var buf bytes.Buffer
	if err := format.Node(&buf, fileSet, f); err != nil {
		return "", err
	}

	return buf.String(), nil
}

func deleteImportsByPath(fileSet *token.FileSet, file *ast.File, path string) {
	names := make([]string, 0, len(file.Imports))
	for _, spec := range file.Imports {
		if importPath(spec) == path {
			names = append(names, importName(spec))
		}
	}

	for _, name := range names {
		astutil.DeleteNamedImport(fileSet, file, name, path)
	}
}

func hasImport(file *ast.File, name, path string) bool {
	for _, spec := range file.Imports {
		if importName(spec) == name && importPath(spec) == path {
			return true
		}
	}
	return false
}

func importName(spec *ast.ImportSpec) string {
	if spec.Name == nil {
		return ""
	}
	return spec.Name.Name
}

func importPath(spec *ast.ImportSpec) string {
	value, err := strconv.Unquote(spec.Path.Value)
	if err != nil {
		return ""
	}
	return value
}