gocomplete: run only runnable go files

when typing 'go run', the completion will complete only go files
which are in main package and have a main function.

5 files changed, 169 insertions, 49 deletions

diff --git a/gocomplete/complete.go b/gocomplete/complete.go
index 5c24417..f90b12b 100644
--- a/gocomplete/complete.go
+++ b/gocomplete/complete.go
@@ -5,7 +5,7 @@ import "github.com/posener/complete"
 
 var (
 	ellipsis   = complete.PredictSet("./...")
-	anyPackage = predictPackages("")
+	anyPackage = complete.PredictFunc(predictPackages)
 	goFiles    = complete.PredictFiles("*.go")
 	anyFile    = complete.PredictFiles("*")
 	anyGo      = complete.PredictOr(goFiles, anyPackage, ellipsis)
@@ -44,7 +44,7 @@ func main() {
 		Flags: complete.Flags{
 			"-exec": complete.PredictAnything,
 		},
-		Args: goFiles,
+		Args: complete.PredictFunc(predictRunnableFiles),
 	}
 
 	test := complete.Command{
@@ -53,14 +53,14 @@ func main() {
 			"-c":    complete.PredictNothing,
 			"-exec": complete.PredictAnything,
 
-			"-bench":     predictTest("Benchmark"),
+			"-bench":     predictBenchmark,
 			"-benchtime": complete.PredictAnything,
 			"-count":     complete.PredictAnything,
 			"-cover":     complete.PredictNothing,
 			"-covermode": complete.PredictSet("set", "count", "atomic"),
 			"-coverpkg":  complete.PredictDirs("*"),
 			"-cpu":       complete.PredictAnything,
-			"-run":       predictTest("Test", "Example"),
+			"-run":       predictTest,
 			"-short":     complete.PredictNothing,
 			"-timeout":   complete.PredictAnything,
 
diff --git a/gocomplete/parse.go b/gocomplete/parse.go
new file mode 100644
index 0000000..8111b74
--- /dev/null
+++ b/gocomplete/parse.go
@@ -0,0 +1,28 @@
+package main
+
+import (
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"regexp"
+
+	"github.com/posener/complete"
+)
+
+func functionsInFile(path string, regexp *regexp.Regexp) (tests []string) {
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, path, nil, 0)
+	if err != nil {
+		complete.Log("Failed parsing %s: %s", path, err)
+		return nil
+	}
+	for _, d := range f.Decls {
+		if f, ok := d.(*ast.FuncDecl); ok {
+			name := f.Name.String()
+			if regexp == nil || regexp.MatchString(name) {
+				tests = append(tests, name)
+			}
+		}
+	}
+	return
+}
diff --git a/gocomplete/pkgs.go b/gocomplete/pkgs.go
index b223ea9..356a0fa 100644
--- a/gocomplete/pkgs.go
+++ b/gocomplete/pkgs.go
@@ -4,37 +4,59 @@ import (
 	"bytes"
 	"encoding/json"
 	"os/exec"
+	"path/filepath"
+	"regexp"
 	"strings"
 
 	"github.com/posener/complete"
 )
 
-const goListFormat = `'{"name": "{{.Name}}", "dir": "{{.Dir}}"}'`
+const goListFormat = `{"Name": "{{.Name}}", "Path": "{{.Dir}}", "FilesString": "{{.GoFiles}}"}`
 
-func predictPackages(packageName string) complete.Predictor {
-	return complete.PredictFunc(func(a complete.Args) (prediction []string) {
-		dir := a.Directory()
-		dir = strings.TrimRight(dir, "/.") + "/..."
+// regexp matches a main function
+var reMainFunc = regexp.MustCompile("^main$")
 
-		pkgs := listPackages(dir)
+func predictPackages(a complete.Args) (prediction []string) {
+	dir := a.Directory()
+	pkgs := listPackages(dir)
 
-		files := make([]string, 0, len(pkgs))
-		for _, p := range pkgs {
-			if packageName != "" && p.Name != packageName {
-				continue
+	files := make([]string, 0, len(pkgs))
+	for _, p := range pkgs {
+		files = append(files, p.Path)
+	}
+	return complete.PredictFilesSet(files).Predict(a)
+}
+
+func predictRunnableFiles(a complete.Args) (prediction []string) {
+	dir := a.Directory()
+	pkgs := listPackages(dir)
+
+	files := []string{}
+	for _, p := range pkgs {
+		// filter non main pacakges
+		if p.Name != "main" {
+			continue
+		}
+		for _, f := range p.Files {
+			path := filepath.Join(p.Path, f)
+			if len(functionsInFile(path, reMainFunc)) > 0 {
+				files = append(files, path)
 			}
-			files = append(files, p.Path)
 		}
-		return complete.PredictFilesSet(files).Predict(a)
-	})
+	}
+	complete.Log("FILES: %s", files)
+	return complete.PredictFilesSet(files).Predict(a)
 }
 
 type pack struct {
-	Name string
-	Path string
+	Name        string
+	Path        string
+	FilesString string
+	Files       []string
 }
 
 func listPackages(dir string) (pkgs []pack) {
+	dir = strings.TrimRight(dir, "/") + "/..."
 	out, err := exec.Command("go", "list", "-f", goListFormat, dir).Output()
 	if err != nil {
 		return
@@ -42,9 +64,13 @@ func listPackages(dir string) (pkgs []pack) {
 	lines := bytes.Split(out, []byte("\n"))
 	for _, line := range lines {
 		var p pack
-		if err := json.Unmarshal(line, &p); err == nil {
-			pkgs = append(pkgs, p)
+		err := json.Unmarshal(line, &p)
+		if err != nil {
+			continue
 		}
+		// parse the FileString from a string "[file1 file2 file3]" to a list of files
+		p.Files = strings.Split(strings.Trim(p.FilesString, "[]"), " ")
+		pkgs = append(pkgs, p)
 	}
 	return
 }
diff --git a/gocomplete/tests.go b/gocomplete/tests.go
index d2c32e7..a952dab 100644
--- a/gocomplete/tests.go
+++ b/gocomplete/tests.go
@@ -1,25 +1,28 @@
 package main
 
 import (
-	"go/ast"
-	"go/parser"
-	"go/token"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 
 	"github.com/posener/complete"
 	"github.com/posener/complete/match"
 )
 
+var (
+	predictBenchmark = funcPredict(regexp.MustCompile("^Benchmark"))
+	predictTest      = funcPredict(regexp.MustCompile("^(Test|Example)"))
+)
+
 // predictTest predict test names.
 // it searches in the current directory for all the go test files
 // and then all the relevant function names.
 // for test names use prefix of 'Test' or 'Example', and for benchmark
 // test names use 'Benchmark'
-func predictTest(funcPrefix ...string) complete.Predictor {
+func funcPredict(funcRegexp *regexp.Regexp) complete.Predictor {
 	return complete.PredictFunc(func(a complete.Args) (prediction []string) {
-		tests := testNames(funcPrefix)
+		tests := funcNames(funcRegexp)
 		for _, t := range tests {
 			if match.Prefix(t, a.Last) {
 				prediction = append(prediction, t)
@@ -30,36 +33,15 @@ func predictTest(funcPrefix ...string) complete.Predictor {
 }
 
 // get all test names in current directory
-func testNames(funcPrefix []string) (tests []string) {
+func funcNames(funcRegexp *regexp.Regexp) (tests []string) {
 	filepath.Walk("./", func(path string, info os.FileInfo, err error) error {
 		// if not a test file, skip
 		if !strings.HasSuffix(path, "_test.go") {
 			return nil
 		}
 		// inspect test file and append all the test names
-		tests = append(tests, testsInFile(funcPrefix, path)...)
+		tests = append(tests, functionsInFile(path, funcRegexp)...)
 		return nil
 	})
 	return
 }
-
-func testsInFile(funcPrefix []string, path string) (tests []string) {
-	fset := token.NewFileSet()
-	f, err := parser.ParseFile(fset, path, nil, 0)
-	if err != nil {
-		complete.Log("Failed parsing %s: %s", path, err)
-		return nil
-	}
-	for _, d := range f.Decls {
-		if f, ok := d.(*ast.FuncDecl); ok {
-			name := f.Name.String()
-			for _, prefix := range funcPrefix {
-				if strings.HasPrefix(name, prefix) {
-					tests = append(tests, name)
-					break
-				}
-			}
-		}
-	}
-	return
-}
diff --git a/gocomplete/tests_test.go b/gocomplete/tests_test.go
new file mode 100644
index 0000000..4eb2308
--- /dev/null
+++ b/gocomplete/tests_test.go
@@ -0,0 +1,84 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/posener/complete"
+)
+
+func TestPredictions(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		predictor  complete.Predictor
+		last       string
+		completion []string
+	}{
+		{
+			name:       "predict tests ok",
+			predictor:  predictTest,
+			completion: []string{"TestPredictions", "Example"},
+		},
+		{
+			name:      "predict tests not found",
+			predictor: predictTest,
+			last:      "X",
+		},
+		{
+			name:       "predict benchmark ok",
+			predictor:  predictBenchmark,
+			completion: []string{"BenchmarkFake"},
+		},
+		{
+			name:      "predict benchmarks not found",
+			predictor: predictBenchmark,
+			last:      "X",
+		},
+		{
+			name:       "predict packages ok",
+			predictor:  complete.PredictFunc(predictPackages),
+			completion: []string{"./"},
+		},
+		{
+			name:      "predict packages not found",
+			predictor: complete.PredictFunc(predictPackages),
+			last:      "X",
+		},
+		{
+			name:       "predict runnable ok",
+			predictor:  complete.PredictFunc(predictRunnableFiles),
+			completion: []string{"./complete.go"},
+		},
+		{
+			name:      "predict runnable not found",
+			predictor: complete.PredictFunc(predictRunnableFiles),
+			last:      "X",
+		},
+	}
+
+	for _, tt := range tests {
+		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)
+			}
+		})
+	}
+}
+
+func BenchmarkFake(b *testing.B) {}
+func Example()                   {}
+
+func equal(s1, s2 []string) bool {
+	if len(s1) != len(s2) {
+		return false
+	}
+	for i := range s1 {
+		if s1[i] != s2[i] {
+			return false
+		}
+	}
+	return true
+}