Add `walkPattern` to return both files and directories

Also added `walkDirs` to walk only directories based on the pattern
specified.

1 files changed, 46 insertions, 10 deletions

diff --git a/lib/pure/os.nim b/lib/pure/os.nim
index 668793c20..52ba110a9 100644
--- a/lib/pure/os.nim
+++ b/lib/pure/os.nim
@@ -778,12 +778,26 @@ iterator envPairs*(): tuple[key, value: TaintedString] {.tags: [ReadEnvEffect].}
     yield (TaintedString(substr(environment[i], 0, p-1)),
            TaintedString(substr(environment[i], p+1)))
 
-iterator walkFiles*(pattern: string): string {.tags: [ReadDirEffect].} =
-  ## Iterate over all the files that match the `pattern`. On POSIX this uses
-  ## the `glob`:idx: call.
-  ##
-  ## `pattern` is OS dependent, but at least the "\*.ext"
-  ## notation is supported.
+# Templates for filtering directories and files
+when defined(windows):
+  template isDir(f: WIN32_FIND_DATA): bool =
+    (f.dwFileAttributes and FILE_ATTRIBUTE_DIRECTORY) != 0'i32
+  template isFile(f: WIN32_FIND_DATA): bool =
+    not isDir(f)
+else:
+  template isDir(f: string): bool =
+    dirExists(f)
+  template isFile(f: string): bool =
+    fileExists(f)
+
+template defaultWalkFilter(item): bool =
+  ## Walk filter used to return true on both
+  ## files and directories
+  true
+
+template walkCommon(pattern: string, filter) =
+  ## Common code for getting the files and directories with the
+  ## specified `pattern`
   when defined(windows):
     var
       f: WIN32_FIND_DATA
@@ -792,8 +806,7 @@ iterator walkFiles*(pattern: string): string {.tags: [ReadDirEffect].} =
     if res != -1:
       defer: findClose(res)
       while true:
-        if not skipFindData(f) and
-            (f.dwFileAttributes and FILE_ATTRIBUTE_DIRECTORY) == 0'i32:
+        if not skipFindData(f) and filter(f):
           # Windows bug/gotcha: 't*.nim' matches 'tfoo.nims' -.- so we check
           # that the file extensions have the same length ...
           let ff = getFilename(f)
@@ -816,10 +829,33 @@ iterator walkFiles*(pattern: string): string {.tags: [ReadDirEffect].} =
       for i in 0.. f.gl_pathc - 1:
         assert(f.gl_pathv[i] != nil)
         let path = $f.gl_pathv[i]
-        # Make sure it's a file and not a directory
-        if fileExists(path):
+        if filter(path):
           yield path
 
+iterator walkPattern*(pattern: string): string {.tags: [ReadDirEffect].} =
+  ## Iterate over all the files and directories that match the `pattern`.
+  ## On POSIX this uses the `glob`:idx: call.
+  ##
+  ## `pattern` is OS dependent, but at least the "\*.ext"
+  ## notation is supported.
+  walkCommon(pattern, defaultWalkFilter)
+
+iterator walkFiles*(pattern: string): string {.tags: [ReadDirEffect].} =
+  ## Iterate over all the files that match the `pattern`. On POSIX this uses
+  ## the `glob`:idx: call.
+  ##
+  ## `pattern` is OS dependent, but at least the "\*.ext"
+  ## notation is supported.
+  walkCommon(pattern, isFile)
+
+iterator walkDirs*(pattern: string): string {.tags: [ReadDirEffect].} =
+  ## Iterate over all the directories that match the `pattern`.
+  ## On POSIX this uses the `glob`:idx: call.
+  ##
+  ## `pattern` is OS dependent, but at least the "\*.ext"
+  ## notation is supported.
+  walkCommon(pattern, isDir)
+
 type
   PathComponent* = enum   ## Enumeration specifying a path component.
     pcFile,               ## path refers to a file