aefix: Drop "appengine" (or updated equivalent) import if it becomes unused.

Change-Id: I879d81a24768afea7598964132204ac9ed2fea78

cmd/aefix/ae.go

@@ -52,11 +52,17 @@ func aeFn(f *ast.File) bool {
 	fixed := false
 
 	// Update imports.
+	mainImp := "appengine"
 	for _, imp := range f.Imports {
 		pth, _ := strconv.Unquote(imp.Path.Value)
 		if pth == "appengine" || strings.HasPrefix(pth, "appengine/") {
-			imp.Path.Value = strconv.Quote(mapPackage(pth))
+			newPth := mapPackage(pth)
+			imp.Path.Value = strconv.Quote(newPth)
 			fixed = true
+
+			if pth == "appengine" {
+				mainImp = newPth
+			}
 		}
 	}
 
@@ -145,8 +151,6 @@ func aeFn(f *ast.File) bool {
 	// Change any `appengine.Context` to `context.Context`.
 	// Do this in a separate walk because the previous walk
 	// wants to identify "appengine.Context".
-	//
-	// TODO(dsymonds): Drop the "appengine" import if it is now orphaned.
 	walk(f, func(n interface{}) {
 		expr, ok := n.(ast.Expr)
 		if ok && isPkgDot(expr, "appengine", "Context") {
@@ -158,6 +162,12 @@ func aeFn(f *ast.File) bool {
 		}
 	})
 
+	// The changes above might remove the need to import "appengine".
+	// Check if it's used, and drop it if it isn't.
+	if fixed && !usesImport(f, mainImp) {
+		deleteImport(f, mainImp)
+	}
+
 	return fixed
 }
 

cmd/aefix/ae_test.go

@@ -75,7 +75,6 @@ func LogSomething(c2 appengine.Context) {
 
 import (
 	"golang.org/x/net/context"
-	"google.golang.org/appengine"
 	"google.golang.org/appengine/log"
 )
 
@@ -104,13 +103,42 @@ func f(ctx appengine.Context) {
 
 import (
 	"golang.org/x/net/context"
-	"google.golang.org/appengine"
 	"google.golang.org/appengine/taskqueue"
 )
 
 func f(ctx context.Context) {
 	stats, err := taskqueue.QueueStats(ctx, []string{"one", "two"})
 }
+`,
+	},
+
+	// Check that the main "appengine" import will not be dropped
+	// if an appengine.Context -> context.Context change happens
+	// but the appengine package is still referenced.
+	{
+		Name: "ae.3",
+		In: `package foo
+
+import (
+	"appengine"
+	"io"
+)
+
+func f(ctx appengine.Context, w io.Writer) {
+	_ = appengine.IsDevAppServer()
+}
+`,
+		Out: `package foo
+
+import (
+	"golang.org/x/net/context"
+	"google.golang.org/appengine"
+	"io"
+)
+
+func f(ctx context.Context, w io.Writer) {
+	_ = appengine.IsDevAppServer()
+}
 `,
 	},
 }