diff --git a/astutil/imports.go b/astutil/imports.go index 9ce61248..c3128103 100644 --- a/astutil/imports.go +++ b/astutil/imports.go @@ -11,6 +11,17 @@ import ( // AddImport adds the import path to the file f, if absent. func AddImport(f *ast.File, ipath string) (added bool) { + return AddNamedImport(f, "", ipath) +} + +// AddNamedImport adds the import path to the file f, if absent. +// If name is not empty, it is used to rename the import. +// +// For example, calling +// AddNamedImport(f, "pathpkg", "path") +// adds +// import pathpkg "path" +func AddNamedImport(f *ast.File, name, ipath string) (added bool) { if imports(f, ipath) { return false } @@ -21,6 +32,9 @@ func AddImport(f *ast.File, ipath string) (added bool) { Value: strconv.Quote(ipath), }, } + if name != "" { + newImport.Name = &ast.Ident{Name: name} + } // Find an import decl to add to. var ( diff --git a/astutil/imports_test.go b/astutil/imports_test.go index 3ecfe518..b0bb00e1 100644 --- a/astutil/imports_test.go +++ b/astutil/imports_test.go @@ -30,10 +30,11 @@ func print(t *testing.T, name string, f *ast.File) string { } type test struct { - name string - pkg string - in string - out string + name string + renamedPkg string + pkg string + in string + out string } var addTests = []test{ @@ -150,6 +151,23 @@ import ( "bytes" "os" ) +`, + }, + { + name: "", + renamedPkg: "fmtpkg", + pkg: "fmt", + in: `package main + +import "os" + +`, + out: `package main + +import ( + fmtpkg "fmt" + "os" +) `, }, } @@ -157,7 +175,7 @@ import ( func TestAddImport(t *testing.T) { for _, test := range addTests { file := parse(t, test.name, test.in) - AddImport(file, test.pkg) + AddNamedImport(file, test.renamedPkg, test.pkg) if got := print(t, test.name, file); got != test.out { t.Errorf("%s:\ngot: %s\nwant: %s", test.name, got, test.out) }