mirror of https://github.com/stashapp/stash.git
53 lines
1.1 KiB
Go
53 lines
1.1 KiB
Go
package gogen
|
|
|
|
import (
|
|
"fmt"
|
|
"go/ast"
|
|
"go/token"
|
|
"strings"
|
|
|
|
"github.com/gobuffalo/genny"
|
|
"github.com/pkg/errors"
|
|
)
|
|
|
|
// AddImport adds n number of import statements into the path provided
|
|
func AddImport(gf genny.File, imports ...string) (genny.File, error) {
|
|
pf, err := ParseFile(gf)
|
|
if err != nil {
|
|
return gf, errors.WithStack(err)
|
|
}
|
|
gf = pf.File
|
|
|
|
end := findLastImport(pf.Ast, pf.FileSet, pf.Lines)
|
|
|
|
x := make([]string, len(imports), len(imports)+2)
|
|
for _, i := range imports {
|
|
x = append(x, fmt.Sprintf("\t\"%s\"", i))
|
|
|
|
}
|
|
if end < 0 {
|
|
x = append([]string{"import ("}, x...)
|
|
x = append(x, ")")
|
|
}
|
|
|
|
pf.Lines = append(pf.Lines[:end], append(x, pf.Lines[end:]...)...)
|
|
|
|
fileContent := strings.Join(pf.Lines, "\n")
|
|
return genny.NewFile(gf.Name(), strings.NewReader(fileContent)), nil
|
|
}
|
|
|
|
func findLastImport(f *ast.File, fset *token.FileSet, fileLines []string) int {
|
|
var end = -1
|
|
|
|
ast.Inspect(f, func(n ast.Node) bool {
|
|
switch x := n.(type) {
|
|
case *ast.ImportSpec:
|
|
end = fset.Position(x.End()).Line
|
|
return true
|
|
}
|
|
return true
|
|
})
|
|
|
|
return end
|
|
}
|