1 files changed, 43 insertions, 49 deletions

diff --git a/predict_files.go b/predict_files.go
index fe70c97..0543442 100644
--- a/predict_files.go
+++ b/predict_files.go
@@ -30,10 +30,9 @@ func files(pattern string, allowFiles bool) PredictFunc {
 	// if only one directory has matched the result, search recursively into
 	// this directory to give more results.
 	return func(a Args) (prediction []string) {
-		last := a.Last
 		for {
 
-			prediction = predictFiles(last, pattern, allowFiles)
+			prediction = predictFiles(a, pattern, allowFiles)
 
 			// if the number of prediction is not 1, we either have many results or
 			// have no results, so we return it.
@@ -43,7 +42,7 @@ func files(pattern string, allowFiles bool) PredictFunc {
 
 			// if the result is only one item, we might want to recursively check
 			// for more accurate results.
-			if prediction[0] == last { // avoid loop forever
+			if prediction[0] == a.Last { // avoid loop forever
 				return
 			}
 
@@ -52,56 +51,59 @@ func files(pattern string, allowFiles bool) PredictFunc {
 				return
 			}
 
-			last = prediction[0]
+			a.Last = prediction[0]
 		}
 	}
 }
 
-func predictFiles(last string, pattern string, allowFiles bool) (prediction []string) {
-	if strings.HasSuffix(last, "/..") {
-		return
+func predictFiles(a Args, pattern string, allowFiles bool) []string {
+	if strings.HasSuffix(a.Last, "/..") {
+		return nil
 	}
 
-	dir := dirFromLast(last)
-	rel := !filepath.IsAbs(pattern)
-	files := listFiles(dir, pattern)
-
-	// get wording directory for relative name
-	workDir, err := os.Getwd()
-	if err != nil {
-		workDir = ""
-	}
+	dir := a.Directory()
+	files := listFiles(dir, pattern, allowFiles)
 
 	// 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
-		}
+	return PredictFilesSet(files).Predict(a)
+}
 
-		// change file name to relative if necessary
-		if rel && workDir != "" {
-			f = toRel(workDir, f)
-		}
+// PredictFilesSet predict according to file rules to a given set of file names
+func PredictFilesSet(files []string) PredictFunc {
+	return func(a Args) (prediction []string) {
+		rel := !filepath.IsAbs(a.Directory())
+		// add all matching files to prediction
+		for _, f := range files {
+			// change file name to relative if necessary
+			if rel {
+				f = toRel(f)
+			}
 
-		// test matching of file to the argument
-		if match.File(f, last) {
-			prediction = append(prediction, f)
+			// test matching of file to the argument
+			if match.File(f, a.Last) {
+				prediction = append(prediction, f)
+			}
 		}
+		return
 	}
-	return
-
 }
 
-func listFiles(dir, pattern string) []string {
+func listFiles(dir, pattern string, allowFiles bool) []string {
+	// set of all file names
 	m := map[string]bool{}
+
+	// list files
 	if files, err := filepath.Glob(filepath.Join(dir, pattern)); err == nil {
 		for _, f := range files {
-			m[f] = true
+			if stat, err := os.Stat(f); err != nil || stat.IsDir() || allowFiles {
+				m[f] = true
+			}
 		}
 	}
+
+	// list directories
 	if dirs, err := ioutil.ReadDir(dir); err == nil {
 		for _, d := range dirs {
 			if d.IsDir() {
@@ -109,6 +111,7 @@ func listFiles(dir, pattern string) []string {
 			}
 		}
 	}
+
 	list := make([]string, 0, len(m))
 	for k := range m {
 		list = append(list, k)
@@ -117,12 +120,18 @@ func listFiles(dir, pattern string) []string {
 }
 
 // toRel changes a file name to a relative name
-func toRel(wd, file string) string {
+func toRel(file string) string {
+	// get wording directory for relative name
+	workDir, err := os.Getwd()
+	if err != nil {
+		return file
+	}
+
 	abs, err := filepath.Abs(file)
 	if err != nil {
 		return file
 	}
-	rel, err := filepath.Rel(wd, abs)
+	rel, err := filepath.Rel(workDir, abs)
 	if err != nil {
 		return file
 	}
@@ -134,18 +143,3 @@ func toRel(wd, file string) string {
 	}
 	return rel
 }
-
-// dirFromLast gives the directory of the current written
-// last argument if it represents a file name being written.
-// in case that it is not, we fall back to the current directory.
-func dirFromLast(last string) string {
-	if info, err := os.Stat(last); err == nil && info.IsDir() {
-		return last
-	}
-	dir := filepath.Dir(last)
-	_, err := os.Stat(dir)
-	if err != nil {
-		return "./"
-	}
-	return dir
-}