bpo-28682: Added support for bytes paths in os.fwalk(). (python#489)

serhiy-storchaka · web-flow · commit 8f6b344d368c · 2017-03-07T14:33:21.000+02:00
diff --git a/Doc/library/os.rst b/Doc/library/os.rst
@@ -2840,6 +2840,9 @@ features:
    .. versionchanged:: 3.6
       Accepts a :term:`path-like object`.
 
+   .. versionchanged:: 3.7
+      Added support for :class:`bytes` paths.
+
 
 Linux extended attributes
 ~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/Doc/whatsnew/3.7.rst b/Doc/whatsnew/3.7.rst
@@ -97,6 +97,12 @@ New Modules
 Improved Modules
 ================
 
+os
+--
+
+Added support for :class:`bytes` paths in :func:`~os.fwalk`.
+(Contributed by Serhiy Storchaka in :issue:`28682`.)
+
 unittest.mock
 -------------
 
diff --git a/Lib/os.py b/Lib/os.py
@@ -460,16 +460,19 @@ def fwalk(top=".", topdown=True, onerror=None, *, follow_symlinks=False, dir_fd=
         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, isinstance(top, bytes),
+                                  topdown, onerror, follow_symlinks)
         finally:
             close(topfd)
 
-    def _fwalk(topfd, toppath, topdown, onerror, follow_symlinks):
+    def _fwalk(topfd, toppath, isbytes, topdown, onerror, follow_symlinks):
         # 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.
 
         names = listdir(topfd)
+        if isbytes:
+            names = map(fsencode, names)
         dirs, nondirs = [], []
         for name in names:
             try:
@@ -504,7 +507,8 @@ def _fwalk(topfd, toppath, topdown, onerror, follow_symlinks):
             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, isbytes,
+                                      topdown, onerror, follow_symlinks)
             finally:
                 close(dirfd)
 
diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
@@ -1010,9 +1010,12 @@ class FwalkTests(WalkTests):
     """Tests for os.fwalk()."""
 
     def walk(self, top, **kwargs):
-        for root, dirs, files, root_fd in os.fwalk(top, **kwargs):
+        for root, dirs, files, root_fd in self.fwalk(top, **kwargs):
             yield (root, dirs, files)
 
+    def fwalk(self, *args, **kwargs):
+        return os.fwalk(*args, **kwargs)
+
     def _compare_to_walk(self, walk_kwargs, fwalk_kwargs):
         """
         compare with walk() results.
@@ -1027,7 +1030,7 @@ def _compare_to_walk(self, walk_kwargs, fwalk_kwargs):
             for root, dirs, files in os.walk(**walk_kwargs):
                 expected[root] = (set(dirs), set(files))
 
-            for root, dirs, files, rootfd in os.fwalk(**fwalk_kwargs):
+            for root, dirs, files, rootfd in self.fwalk(**fwalk_kwargs):
                 self.assertIn(root, expected)
                 self.assertEqual(expected[root], (set(dirs), set(files)))
 
@@ -1049,7 +1052,7 @@ def test_yields_correct_dir_fd(self):
         # check returned file descriptors
         for topdown, follow_symlinks in itertools.product((True, False), repeat=2):
             args = support.TESTFN, topdown, None
-            for root, dirs, files, rootfd in os.fwalk(*args, follow_symlinks=follow_symlinks):
+            for root, dirs, files, rootfd in self.fwalk(*args, follow_symlinks=follow_symlinks):
                 # check that the FD is valid
                 os.fstat(rootfd)
                 # redundant check
@@ -1064,22 +1067,14 @@ def test_fd_leak(self):
         minfd = os.dup(1)
         os.close(minfd)
         for i in range(256):
-            for x in os.fwalk(support.TESTFN):
+            for x in self.fwalk(support.TESTFN):
                 pass
         newfd = os.dup(1)
         self.addCleanup(os.close, newfd)
         self.assertEqual(newfd, minfd)
 
 class BytesWalkTests(WalkTests):
     """Tests for os.walk() with bytes."""
-    def setUp(self):
-        super().setUp()
-        self.stack = contextlib.ExitStack()
-
-    def tearDown(self):
-        self.stack.close()
-        super().tearDown()
-
     def walk(self, top, **kwargs):
         if 'follow_symlinks' in kwargs:
             kwargs['followlinks'] = kwargs.pop('follow_symlinks')
@@ -1091,6 +1086,18 @@ def walk(self, top, **kwargs):
             bdirs[:] = list(map(os.fsencode, dirs))
             bfiles[:] = list(map(os.fsencode, files))
 
+@unittest.skipUnless(hasattr(os, 'fwalk'), "Test needs os.fwalk()")
+class BytesFwalkTests(FwalkTests):
+    """Tests for os.walk() with bytes."""
+    def fwalk(self, top='.', *args, **kwargs):
+        for broot, bdirs, bfiles, topfd in os.fwalk(os.fsencode(top), *args, **kwargs):
+            root = os.fsdecode(broot)
+            dirs = list(map(os.fsdecode, bdirs))
+            files = list(map(os.fsdecode, bfiles))
+            yield (root, dirs, files, topfd)
+            bdirs[:] = list(map(os.fsencode, dirs))
+            bfiles[:] = list(map(os.fsencode, files))
+
 
 class MakedirTests(unittest.TestCase):
     def setUp(self):
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -265,6 +265,8 @@ Extension Modules
 Library
 -------
 
+- bpo-28682: Added support for bytes paths in os.fwalk().
+
 - bpo-29623: Allow use of path-like object as a single argument in
   ConfigParser.read().  Patch by David Ellis.
 
