1 files changed, 54 insertions, 50 deletions

diff --git a/predict.go b/predict.go
index 1db33f5..64e793b 100644
--- a/predict.go
+++ b/predict.go
@@ -1,6 +1,7 @@
 package complete
 
 import (
+	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strings"
@@ -68,7 +69,7 @@ func (p predictSet) Predict(a Args) (prediction []string) {
 // path, if no path was started to be typed, it will complete to directories
 // in the current working directory.
 func PredictDirs(pattern string) Predictor {
-	return files(pattern, true, false)
+	return files(pattern, false)
 }
 
 // PredictFiles will search for files matching the given pattern in the started to
@@ -76,34 +77,40 @@ func PredictDirs(pattern string) Predictor {
 // match the pattern in the current working directory.
 // To match any file, use "*" as pattern. To match go files use "*.go", and so on.
 func PredictFiles(pattern string) Predictor {
-	return files(pattern, false, true)
+	return files(pattern, true)
 }
 
-// PredictFilesOrDirs any file or directory that matches the pattern
-func PredictFilesOrDirs(pattern string) Predictor {
-	return files(pattern, true, true)
-}
-
-func files(pattern string, allowDirs, allowFiles bool) PredictFunc {
+func files(pattern string, allowFiles bool) PredictFunc {
 	return func(a Args) (prediction []string) {
 		if strings.HasSuffix(a.Last, "/..") {
 			return
 		}
 		dir := dirFromLast(a.Last)
+		rel := !filepath.IsAbs(pattern)
 		Log("looking for files in %s (last=%s)", dir, a.Last)
-		files, err := filepath.Glob(filepath.Join(dir, pattern))
+		files := listFiles(dir, pattern)
+
+		// get wording directory for relative name
+		workDir, err := os.Getwd()
 		if err != nil {
-			Log("failed glob operation with pattern '%s': %s", pattern, err)
-		}
-		if allowDirs {
-			files = append(files, dir)
-		}
-		files = selectByType(files, allowDirs, allowFiles)
-		if !filepath.IsAbs(pattern) {
-			filesToRel(files)
+			workDir = ""
 		}
+
+		// add dir if match
+		files = append(files, dir)
+
 		// add all matching files to prediction
 		for _, f := range files {
+			if stat, err := os.Stat(f); err != nil || (!stat.IsDir() && !allowFiles) {
+				continue
+			}
+
+			// change file name to relative if necessary
+			if rel && workDir != "" {
+				f = toRel(workDir, f)
+			}
+
+			// test matching of file to the argument
 			if match.File(f, a.Last) {
 				prediction = append(prediction, f)
 			}
@@ -111,47 +118,44 @@ func files(pattern string, allowDirs, allowFiles bool) PredictFunc {
 		return
 	}
 }
-
-func selectByType(names []string, allowDirs bool, allowFiles bool) []string {
-	filtered := make([]string, 0, len(names))
-	for _, name := range names {
-		stat, err := os.Stat(name)
-		if err != nil {
-			continue
+func listFiles(dir, pattern string) []string {
+	m := map[string]bool{}
+	if files, err := filepath.Glob(filepath.Join(dir, pattern)); err == nil {
+		for _, f := range files {
+			m[f] = true
 		}
-		if (stat.IsDir() && !allowDirs) || (!stat.IsDir() && !allowFiles) {
-			continue
+	}
+	if dirs, err := ioutil.ReadDir(dir); err == nil {
+		for _, d := range dirs {
+			if d.IsDir() {
+				m[d.Name()] = true
+			}
 		}
-		filtered = append(filtered, name)
 	}
-	return filtered
+	list := make([]string, 0, len(m))
+	for k := range m {
+		list = append(list, k)
+	}
+	return list
 }
 
-// filesToRel, change list of files to their names in the relative
-// to current directory form.
-func filesToRel(files []string) {
-	wd, err := os.Getwd()
+// toRel changes a file name to a relative name
+func toRel(wd, file string) string {
+	abs, err := filepath.Abs(file)
 	if err != nil {
-		return
+		return file
 	}
-	for i := range files {
-		abs, err := filepath.Abs(files[i])
-		if err != nil {
-			continue
-		}
-		rel, err := filepath.Rel(wd, abs)
-		if err != nil {
-			continue
-		}
-		if rel != "." {
-			rel = "./" + rel
-		}
-		if info, err := os.Stat(rel); err == nil && info.IsDir() {
-			rel += "/"
-		}
-		files[i] = rel
+	rel, err := filepath.Rel(wd, abs)
+	if err != nil {
+		return file
 	}
-	return
+	if rel != "." {
+		rel = "./" + rel
+	}
+	if info, err := os.Stat(rel); err == nil && info.IsDir() {
+		rel += "/"
+	}
+	return rel
 }
 
 // dirFromLast gives the directory of the current written