From a67bb12457c4c976419760835140490b783ec9f0 Mon Sep 17 00:00:00 2001
From: Eyal Posener <posener@gmail.com>
Date: Fri, 28 Jul 2017 15:01:07 +0300
Subject: [PATCH] gocomplete: Add support for system GOPATH packages

fixes #41
---
 gocomplete/pkgs.go       | 66 +++++++++++++++++++++++++++++++++++++---
 gocomplete/tests_test.go | 37 +++++++++++++---------
 2 files changed, 84 insertions(+), 19 deletions(-)

diff --git a/gocomplete/pkgs.go b/gocomplete/pkgs.go
index d1f0331..8110ff9 100644
--- a/gocomplete/pkgs.go
+++ b/gocomplete/pkgs.go
@@ -3,7 +3,9 @@ package main
 import (
 	"go/build"
 	"io/ioutil"
+	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/posener/complete"
 )
@@ -11,11 +13,29 @@ import (
 // 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) {
-	prediction = complete.PredictFilesSet(listPackages(a.Directory())).Predict(a)
-	if len(prediction) != 1 {
-		return
+	prediction = []string{a.Last}
+	lastPrediction := ""
+	for len(prediction) == 1 && (lastPrediction == "" || lastPrediction != prediction[0]) {
+		// if only one prediction, predict files within this prediction,
+		// for example, if the user entered 'pk' and we have a package named 'pkg',
+		// which is the only package prefixed with 'pk', we will automatically go one
+		// level deeper and give the user the 'pkg' and all the nested packages within
+		// that package.
+		lastPrediction = prediction[0]
+		a.Last = prediction[0]
+		prediction = predictLocalAndSystem(a)
 	}
-	return complete.PredictFilesSet(listPackages(prediction[0])).Predict(a)
+	return
+}
+
+func predictLocalAndSystem(a complete.Args) []string {
+	localDirs := complete.PredictFilesSet(listPackages(a.Directory())).Predict(a)
+	// System directories are not actual file names, for example: 'github.com/posener/complete' could
+	// be the argument, but the actual filename is in $GOPATH/src/github.com/posener/complete'. this
+	// is the reason to use the PredictSet and not the PredictDirs in this case.
+	s := systemDirs(a.Last)
+	sysDirs := complete.PredictSet(s...).Predict(a)
+	return append(localDirs, sysDirs...)
 }
 
 // listPackages looks in current pointed dir and in all it's direct sub-packages
@@ -48,3 +68,41 @@ func listPackages(dir string) (directories []string) {
 	}
 	return
 }
+
+func systemDirs(dir string) (directories []string) {
+	// get all paths from GOPATH environment variable and use their src directory
+	paths := strings.Split(os.Getenv("GOPATH"), ":")
+	for i := range paths {
+		paths[i] = filepath.Join(paths[i], "src")
+	}
+
+	// normalize the directory to be an actual directory since it could be with an additional
+	// characters after the last '/'.
+	if !strings.HasSuffix(dir, "/") {
+		dir = filepath.Dir(dir)
+	}
+
+	for _, basePath := range paths {
+		path := filepath.Join(basePath, dir)
+		files, err := ioutil.ReadDir(path)
+		if err != nil {
+			// path does not exists
+			continue
+		}
+		// add the base path as one of the completion options
+		switch dir {
+		case "", ".", "/", "./":
+		default:
+			directories = append(directories, dir)
+		}
+		// add all nested directories of the base path
+		// go supports only packages and not go files within the GOPATH
+		for _, f := range files {
+			if !f.IsDir() {
+				continue
+			}
+			directories = append(directories, filepath.Join(dir, f.Name())+"/")
+		}
+	}
+	return
+}
diff --git a/gocomplete/tests_test.go b/gocomplete/tests_test.go
index 5683d24..6799157 100644
--- a/gocomplete/tests_test.go
+++ b/gocomplete/tests_test.go
@@ -12,15 +12,15 @@ func TestPredictions(t *testing.T) {
 	t.Parallel()
 
 	tests := []struct {
-		name       string
-		predictor  complete.Predictor
-		last       string
-		completion []string
+		name      string
+		predictor complete.Predictor
+		last      string
+		want      []string
 	}{
 		{
-			name:       "predict tests ok",
-			predictor:  predictTest,
-			completion: []string{"TestPredictions", "Example"},
+			name:      "predict tests ok",
+			predictor: predictTest,
+			want:      []string{"TestPredictions", "Example"},
 		},
 		{
 			name:      "predict tests not found",
@@ -28,9 +28,9 @@ func TestPredictions(t *testing.T) {
 			last:      "X",
 		},
 		{
-			name:       "predict benchmark ok",
-			predictor:  predictBenchmark,
-			completion: []string{"BenchmarkFake"},
+			name:      "predict benchmark ok",
+			predictor: predictBenchmark,
+			want:      []string{"BenchmarkFake"},
 		},
 		{
 			name:      "predict benchmarks not found",
@@ -38,9 +38,16 @@ func TestPredictions(t *testing.T) {
 			last:      "X",
 		},
 		{
-			name:       "predict packages ok",
-			predictor:  complete.PredictFunc(predictPackages),
-			completion: []string{"./"},
+			name:      "predict local ok",
+			predictor: complete.PredictFunc(predictPackages),
+			last:      ".",
+			want:      []string{"./"},
+		},
+		{
+			name:      "predict system ok",
+			predictor: complete.PredictFunc(predictPackages),
+			last:      "github.com/posener/complete/goc",
+			want:      []string{"github.com/posener/complete/gocomplete/"},
 		},
 		{
 			name:      "predict packages not found",
@@ -53,8 +60,8 @@ func TestPredictions(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			a := complete.Args{Last: tt.last}
 			got := tt.predictor.Predict(a)
-			if want := tt.completion; !equal(got, want) {
-				t.Errorf("Failed %s: completion = %q, want %q", t.Name(), got, want)
+			if !equal(got, tt.want) {
+				t.Errorf("Failed %s: got: %q, want: %q", t.Name(), got, tt.want)
 			}
 		})
 	}