1package gogen
2
3import (
4	"bytes"
5	"go/ast"
6	"go/parser"
7	"go/printer"
8
9	"github.com/gobuffalo/genny"
10	"github.com/pkg/errors"
11	"golang.org/x/tools/go/ast/astutil"
12)
13
14func RewriteImports(gf genny.File, swaps map[string]string) (genny.File, error) {
15	pf, err := ParseFileMode(gf, parser.ParseComments)
16	if err != nil {
17		return gf, errors.WithStack(err)
18	}
19	for key, value := range swaps {
20		if !astutil.DeleteImport(pf.FileSet, pf.Ast, key) {
21			continue
22		}
23
24		astutil.AddImport(pf.FileSet, pf.Ast, value)
25	}
26	ast.SortImports(pf.FileSet, pf.Ast)
27
28	w := &bytes.Buffer{}
29	if err = (&printer.Config{Mode: printer.TabIndent | printer.UseSpaces, Tabwidth: 8}).Fprint(w, pf.FileSet, pf.Ast); err != nil {
30		return gf, errors.WithStack(err)
31	}
32
33	return genny.NewFile(gf.Name(), w), nil
34}
35