storage: add depth parameter to LinkStorage

* Adds a "depth" parameter to the LinkStorage class which allows
  deciding the max depth at which symlinks will be created up to.

pywebpack/storage.py

@@ -11,23 +11,44 @@
 
 from __future__ import absolute_import, print_function
 
-from os import makedirs, remove, symlink, walk
+from os import makedirs, remove, scandir, symlink, walk
 from os.path import dirname, exists, getmtime, islink, join, realpath, relpath
 from shutil import copy
 
 
 def iter_files(folder):
-    """Recursively iterate all files in a given root directory."""
+    """Iterate all files in a given root directory."""
     for root, dirnames, filenames in walk(folder):
         for f in filenames:
             f = join(root, f)
             yield (f, relpath(f, folder))
 
 
+def iter_paths(folder, root=None, depth=None):
+    """Recursively yields paths under a folder up to a maximum depth."""
+    assert depth is None or depth >= 0
+    root = root or folder  # needed to compute the relative name
+
+    if depth is None:  # yield all paths
+        yield from iter_files(folder)
+    elif depth == 0:
+        yield folder, relpath(folder, root)
+    else:
+        for entry in scandir(folder):
+            if entry.is_file():  # always yield files no matter the depth
+                f = join(folder, entry.name)
+                yield f, relpath(f, root)
+            else:
+                yield from iter_paths(
+                    join(folder, entry.name),
+                    root=root, depth=(depth - 1)
+                )
+
+
 class FileStorage(object):
     """Storage class that copies files if source is newer than destination."""
 
-    def __init__(self, srcdir, dstdir):
+    def __init__(self, srcdir, dstdir, **kwargs):
         """Initialize storage."""
         self.srcdir = srcdir
         self.dstdir = dstdir
@@ -60,6 +81,16 @@ def run(self, force=None):
 class LinkStorage(FileStorage):
     """Storage class that link files."""
 
+    def __init__(self, *args, depth=None, **kwargs):
+        """Initialize storage."""
+        self.depth = depth
+        super(LinkStorage, self).__init__(*args, **kwargs)
+
+    def __iter__(self):
+        """Iterate files from a directory."""
+        # Only yield files and directories up to "depth"
+        return iter_paths(self.srcdir, depth=self.depth)
+
     def _copyfile(self, src, dst, force=False):
         """Symlink file from source to destination."""
         if exists(dst):

tests/test_storage.py

@@ -15,7 +15,7 @@
 from os import remove, symlink, utime
 from os.path import exists, getmtime, islink, join, realpath
 
-from pywebpack.storage import FileStorage, LinkStorage, iter_files
+from pywebpack.storage import FileStorage, LinkStorage, iter_files, iter_paths
 
 
 def test_iterfiles(sourcedir):
@@ -26,12 +26,36 @@ def test_iterfiles(sourcedir):
         'buildtpl/webpack.config.js',
         'bundle/index.js',
         'bundle2/main.js',
+        'just-a-file.js',
         'simple/index.js',
         'simple/package.json',
         'simple/webpack.config.js'
     ]
 
 
+def test_iterpaths(sourcedir):
+    """Test paths iteration."""
+    # no args has same behavior as "iter_files" from above
+    assert set(iter_paths(sourcedir)) == set(iter_files(sourcedir))
+    # depth 2 in this case would be also the same...
+    assert set(iter_paths(sourcedir, depth=2)) == set(iter_files(sourcedir))
+    # ...higher depths shouldn't affect the result either
+    assert set(iter_paths(sourcedir, depth=3)) == set(iter_files(sourcedir))
+
+    # depth 0 returns itself
+    assert set(iter_paths(sourcedir, depth=0)) == {(sourcedir, '.')}
+
+    # depth 1 returns first level of files and folders
+    assert {x[1] for x in iter_paths(sourcedir, depth=1)} == {
+        'broken',
+        'buildtpl',
+        'bundle',
+        'bundle2',
+        'simple',
+        'just-a-file.js',
+    }
+
+
 def test_filestorage(sourcedir, tmpdir):
     """Test file storage copy."""
     fsrc = join(sourcedir, 'simple/package.json')