go.tools/astutil: add AddNamedImport

R=bradfitz
CC=golang-dev
https://golang.org/cl/23850044

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 (

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)
 		}