diff -r 3f3b3d4881f6 Doc/library/os.rst --- a/Doc/library/os.rst Wed Apr 13 15:37:23 2016 +0300 +++ b/Doc/library/os.rst Sat Apr 16 17:00:22 2016 +0300 @@ -2528,7 +2528,7 @@ and the *dir_fd*, *follow_symlinks*, and *ns* parameters. -.. function:: walk(top, topdown=True, onerror=None, followlinks=False) +.. function:: walk(top, topdown=True, onerror=None, followlinks=False, max_depth=None) .. index:: single: directory; walking @@ -2573,11 +2573,18 @@ directories. Set *followlinks* to ``True`` to visit directories pointed to by symlinks, on systems that support them. + By default, :func:`walk` will go throw all files in the tree rooted at directory + *top* (including *top* itself). Set *max_depth* to limit the directory depth + you wish to walk. Setting *max_depth* to 1 will walk only *top* itself, + when *max_dpeth* is 2 only *top* and *top* subdirectories will be walked + and so on. + .. note:: Be aware that setting *followlinks* to ``True`` can lead to infinite recursion if a link points to a parent directory of itself. :func:`walk` - does not keep track of the directories it visited already. + does not keep track of the directories it visited already. You may use + *max_depth* to avoid infinite recursion. .. note:: @@ -2617,8 +2624,12 @@ This function now calls :func:`os.scandir` instead of :func:`os.listdir`, making it faster by reducing the number of calls to :func:`os.stat`. - -.. function:: fwalk(top='.', topdown=True, onerror=None, *, follow_symlinks=False, dir_fd=None) + .. versionchanged:: 3.6 + The *max_depth* optional argument allow the user to specify the walk + maximum directory depth. + + +.. function:: fwalk(top='.', topdown=True, onerror=None, *, follow_symlinks=False, dir_fd=None, max_depth=None) .. index:: single: directory; walking diff -r 3f3b3d4881f6 Lib/os.py --- a/Lib/os.py Wed Apr 13 15:37:23 2016 +0300 +++ b/Lib/os.py Sat Apr 16 17:00:22 2016 +0300 @@ -295,7 +295,7 @@ __all__.extend(["makedirs", "removedirs", "renames"]) -def walk(top, topdown=True, onerror=None, followlinks=False): +def walk(top, topdown=True, onerror=None, followlinks=False, max_depth=None): """Directory tree generator. For each directory in the directory tree rooted at top (including top @@ -353,6 +353,8 @@ dirs.remove('CVS') # don't visit CVS directories """ + if max_depth == 0: + return dirs = [] nondirs = [] @@ -417,6 +419,9 @@ if walk_into: walk_dirs.append(entry.path) + if max_depth is not None: + max_depth -= 1 + # Yield before recursion if going top down if topdown: yield top, dirs, nondirs @@ -430,11 +435,11 @@ # the caller can replace the directory entry during the "yield" # above. if followlinks or not islink(new_path): - yield from walk(new_path, topdown, onerror, followlinks) + yield from walk(new_path, topdown, onerror, followlinks, max_depth) else: # Recurse into sub-directories for new_path in walk_dirs: - yield from walk(new_path, topdown, onerror, followlinks) + yield from walk(new_path, topdown, onerror, followlinks, max_depth) # Yield after recursion if going bottom up yield top, dirs, nondirs @@ -503,7 +508,7 @@ if {open, stat} <= supports_dir_fd and {listdir, stat} <= supports_fd: - def fwalk(top=".", topdown=True, onerror=None, *, follow_symlinks=False, dir_fd=None): + def fwalk(top=".", topdown=True, onerror=None, *, follow_symlinks=False, dir_fd=None, max_depth=None): """Directory tree generator. This behaves exactly like walk(), except that it yields a 4-tuple @@ -543,14 +548,16 @@ try: if (follow_symlinks or (st.S_ISDIR(orig_st.st_mode) and path.samestat(orig_st, stat(topfd)))): - yield from _fwalk(topfd, top, topdown, onerror, follow_symlinks) + yield from _fwalk(topfd, top, topdown, onerror, follow_symlinks, max_depth) finally: close(topfd) - def _fwalk(topfd, toppath, topdown, onerror, follow_symlinks): + def _fwalk(topfd, toppath, topdown, onerror, follow_symlinks, max_depth): # Note: This uses O(depth of the directory tree) file descriptors: if # necessary, it can be adapted to only require O(1) FDs, see issue # #13734. + if max_depth == 0: + return names = listdir(topfd) dirs, nondirs = [], [] @@ -573,6 +580,9 @@ except FileNotFoundError: continue + if max_depth is not None: + max_depth -= 1 + if topdown: yield toppath, dirs, nondirs, topfd @@ -587,7 +597,7 @@ try: if follow_symlinks or path.samestat(orig_st, stat(dirfd)): dirpath = path.join(toppath, name) - yield from _fwalk(dirfd, dirpath, topdown, onerror, follow_symlinks) + yield from _fwalk(dirfd, dirpath, topdown, onerror, follow_symlinks, max_depth) finally: close(dirfd) diff -r 3f3b3d4881f6 Lib/test/test_os.py --- a/Lib/test/test_os.py Wed Apr 13 15:37:23 2016 +0300 +++ b/Lib/test/test_os.py Sat Apr 16 17:00:22 2016 +0300 @@ -938,6 +938,18 @@ for dir2 in dirs[1:]: self.assertIn(os.path.join(root, dir2), roots) + def test_walk_max_depth(self): + all = list(self.walk(self.walk_path, max_depth=2)) + self.assertEqual(len(all), 3) + # We can't know which order SUB1 and SUB2 will appear in. + # Not flipped: TESTFN, SUB1, SUB2 + # flipped: TESTFN, SUB2, SUB1 + flipped = all[0][1][0] != "SUB1" + all[0][1].sort() + all[3 - 2 * flipped][-1].sort() + self.assertEqual(all[0], (self.walk_path, ["SUB1", "SUB2"], ["tmp1"])) + self.assertEqual(all[1 + flipped], (self.sub1_path, ["SUB11"], ["tmp2"])) + self.assertEqual(all[2 - flipped], self.sub2_tree) @unittest.skipUnless(hasattr(os, 'fwalk'), "Test needs os.fwalk()") class FwalkTests(WalkTests):