Quick implementation of max_depth for boltons.fileutils.iter_find_files

boltons/fileutils.py

@@ -486,7 +486,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         return
 
 
-def iter_find_files(directory, patterns, ignored=None, include_dirs=False):
+def iter_find_files(directory, patterns, ignored=None, include_dirs=False, max_depth=None):
     """Returns a generator that yields file paths under a *directory*,
     matching *patterns* using `glob`_ syntax (e.g., ``*.txt``). Also
     supports *ignored* patterns.
@@ -500,6 +500,9 @@ def iter_find_files(directory, patterns, ignored=None, include_dirs=False):
             glob-formatted patterns to ignore.
         include_dirs (bool): Whether to include directories that match
            patterns, as well. Defaults to ``False``.
+        max_depth (int): traverse up to this level of subdirectory.
+           I.e., 0 for the specified *directory* only, 1 for *directory* 
+           and one level of subdirectory.
 
     For example, finding Python files in the current directory:
 
@@ -524,7 +527,10 @@ def iter_find_files(directory, patterns, ignored=None, include_dirs=False):
     elif isinstance(ignored, str):
         ignored = [ignored]
     ign_re = re.compile('|'.join([fnmatch.translate(p) for p in ignored]))
+    start_depth = len(directory.split(os.path.sep))
     for root, dirs, files in os.walk(directory):
+        if max_depth is not None and (len(root.split(os.path.sep)) - start_depth) > max_depth:
+            continue
         if include_dirs:
             for basename in dirs:
                 if pats_re.match(basename):

tests/test_fileutils.py

@@ -1,4 +1,10 @@
-from boltons.fileutils import FilePerms
+import os.path
+
+from boltons import fileutils
+from boltons.fileutils import FilePerms, iter_find_files
+
+
+BOLTONS_PATH = os.path.dirname(os.path.abspath(fileutils.__file__))
 
 
 def test_fileperms():
@@ -17,3 +23,14 @@ def test_fileperms():
     assert oct(int(up)) == '0o770'
 
     assert int(FilePerms()) == 0
+
+
+def test_iter_find_files():
+    def _to_baseless_list(paths):
+        return [p.lstrip(BOLTONS_PATH) for p in paths]
+
+    assert 'fileutils.py' in _to_baseless_list(iter_find_files(BOLTONS_PATH, patterns=['*.py']))
+
+    boltons_parent = os.path.dirname(BOLTONS_PATH)
+    assert 'fileutils.py' in _to_baseless_list(iter_find_files(boltons_parent, patterns=['*.py']))
+    assert 'fileutils.py' not in _to_baseless_list(iter_find_files(boltons_parent, patterns=['*.py'], max_depth=0))