gocomplete: improve package completion

gocomplete/pkgs.go

@@ -3,64 +3,48 @@ package main
 import (
 	"go/build"
 	"io/ioutil"
-	"os"
 	"path/filepath"
 
 	"github.com/posener/complete"
 )
 
+// predictPackages completes packages in the directory pointed by a.Last
+// and packages that are one level below that package.
 func predictPackages(a complete.Args) (prediction []string) {
-	for {
-		prediction = complete.PredictFilesSet(listPackages(a)).Predict(a)
-
-		// if the number of prediction is not 1, we either have many results or
-		// have no results, so we return it.
-		if len(prediction) != 1 {
-			return
-		}
-
-		// if the result is only one item, we might want to recursively check
-		// for more accurate results.
-		if prediction[0] == a.Last {
-			return
-		}
-
-		// only try deeper, if the one item is a directory
-		if stat, err := os.Stat(prediction[0]); err != nil || !stat.IsDir() {
-			return
-		}
-
-		a.Last = prediction[0]
-	}
-}
-
-func listPackages(a complete.Args) (dirctories []string) {
-	dir := a.Directory()
-	complete.Log("listing packages in %s", dir)
-	// import current directory
-	pkg, err := build.ImportDir(dir, 0)
-	if err != nil {
-		complete.Log("failed importing directory %s: %s", dir, err)
+	prediction = complete.PredictFilesSet(listPackages(a.Directory())).Predict(a)
+	if len(prediction) != 1 {
 		return
 	}
-	dirctories = append(dirctories, pkg.Dir)
+	return complete.PredictFilesSet(listPackages(prediction[0])).Predict(a)
+}
 
-	// import subdirectories
+// listPackages looks in current pointed dir and in all it's direct sub-packages
+// and return a list of paths to go packages.
+func listPackages(dir string) (directories []string) {
+	// add subdirectories
 	files, err := ioutil.ReadDir(dir)
 	if err != nil {
 		complete.Log("failed reading directory %s: %s", dir, err)
 		return
 	}
+
+	// build paths array
+	paths := make([]string, 0, len(files)+1)
 	for _, f := range files {
-		if !f.IsDir() {
-			continue
+		if f.IsDir() {
+			paths = append(paths, filepath.Join(dir, f.Name()))
 		}
-		pkg, err := build.ImportDir(filepath.Join(dir, f.Name()), 0)
+	}
+	paths = append(paths, dir)
+
+	// import packages according to given paths
+	for _, p := range paths {
+		pkg, err := build.ImportDir(p, 0)
 		if err != nil {
-			complete.Log("failed importing subdirectory %s: %s", filepath.Join(dir, f.Name()), err)
+			complete.Log("failed importing directory %s: %s", p, err)
 			continue
 		}
-		dirctories = append(dirctories, pkg.Dir)
+		directories = append(directories, pkg.Dir)
 	}
 	return
 }