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.

1 files changed, 10 insertions, 28 deletions

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