From e5aae93bfe04225cbad4f4aa3aea943be74f29e0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Oscar=20Nihlg=C3=A5rd?= <oscarnihlgard@gmail.com>
Date: Mon, 26 Nov 2018 13:14:52 +0100
Subject: [PATCH] Add `relative` parameter to walkDirRec

---
 lib/pure/os.nim      | 19 ++++++++++++-------
 tests/stdlib/tos.nim | 19 +++++++++++++++++++
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/lib/pure/os.nim b/lib/pure/os.nim
index e2dd872e8031..0e43e18ca074 100644
--- a/lib/pure/os.nim
+++ b/lib/pure/os.nim
@@ -1564,11 +1564,14 @@ iterator walkDir*(dir: string; relative=false): tuple[kind: PathComponent, path:
               k = getSymlinkFileKind(y)
             yield (k, y)
 
-iterator walkDirRec*(dir: string, yieldFilter = {pcFile},
-                     followFilter = {pcDir}): string {.tags: [ReadDirEffect].} =
+iterator walkDirRec*(dir: string,
+                     yieldFilter = {pcFile}, followFilter = {pcDir},
+                     relative = false): string {.tags: [ReadDirEffect].} =
   ## Recursively walks over the directory `dir` and yields for each file
   ## or directory in `dir`.
-  ## The full path for each file or directory is returned.
+  ## If ``relative`` is true the resulting path is
+  ## shortened to be relative to ``dir``, otherwise the full path is returned.
+  ##
   ## **Warning**:
   ## Modifying the directory structure while the iterator
   ## is traversing may result in undefined behavior!
@@ -1591,13 +1594,15 @@ iterator walkDirRec*(dir: string, yieldFilter = {pcFile},
   ## ``pcLinkToDir``         follow symbolic links to directories
   ## ---------------------   ---------------------------------------------
   ##
-  var stack = @[dir]
+  var stack = @[""]
   while stack.len > 0:
-    for k, p in walkDir(stack.pop()):
+    let d = stack.pop()
+    for k, p in walkDir(dir / d, relative = true):
+      let rel = d / p
       if k in {pcDir, pcLinkToDir} and k in followFilter:
-        stack.add(p)
+        stack.add rel
       if k in yieldFilter:
-        yield p
+        yield if relative: rel else: dir / rel
 
 proc rawRemoveDir(dir: string) {.noNimScript.} =
   when defined(windows):
diff --git a/tests/stdlib/tos.nim b/tests/stdlib/tos.nim
index c10f7036b00a..467f64fffff9 100644
--- a/tests/stdlib/tos.nim
+++ b/tests/stdlib/tos.nim
@@ -170,6 +170,25 @@ block modificationTime:
     echo getLastModificationTime("a") == tm
   removeFile("a")
 
+block walkDirRec:
+  createDir("walkdir_test/a/b")
+  open("walkdir_test/a/b/file_1", fmWrite).close()
+  open("walkdir_test/a/file_2", fmWrite).close()
+
+  for p in walkDirRec("walkdir_test"):
+    doAssert p.fileExists
+    doAssert p.startsWith("walkdir_test")
+
+  var s: seq[string]
+  for p in walkDirRec("walkdir_test", {pcFile}, {pcDir}, relative=true):
+    s.add(p)
+
+  doAssert s.len == 2
+  doAssert "a" / "b" / "file_1" in s
+  doAssert "a" / "file_2" in s
+
+  removeDir("walkdir_test")
+
 block normalizedPath:
   when defined(posix):
     block relative: