Support os.walk(dir_fd=) and os.fwalk(dir_fd=)

Lib/os.py

 r"""OS routines for Mac, NT, or Posix depending on what system we're on.
 
 This exports:
   - all functions from posix, nt, os2, or ce, e.g. unlink, stat, etc.
   - os.path is either posixpath or ntpath
   - os.name is either 'posix', 'nt', 'os2' or 'ce'.
   - os.curdir is a string representing the current directory ('.' or ':')
   - os.pardir is a string representing the parent directory ('..' or '::')
   - os.sep is the (or a most common) pathname separator ('/' or ':' or '\\')
   - os.extsep is the extension separator (always '.')
@@ skipping 313 matching lines @@
     rename(old, new)
     head, tail = path.split(old)
     if head and tail:
         try:
             removedirs(head)
         except error:
             pass
 
 __all__.extend(["makedirs", "removedirs", "renames"])
 
-def walk(top, topdown=True, onerror=None, followlinks=False):
+def _isdir_dir_fd(path, dir_fd):
+    """
+    A substitute for os.path.isdir that supports dir_fd.
+    """
+    try:
+        return st.S_ISDIR(stat(path, dir_fd=dir_fd).st_mode)
+    except FileNotFoundError:
+        return False

[storchaka, 2012/06/25 14:08:10]

Should onerror handler be used?

[larry, 2012/06/25 15:04:13]

I have a much better implementation, so this one is best forgotten.  The new
implementation is just seven lines right at the top of os.walk, and looks
basically like this:

if they pass in dir_fd:
  if fwalk is unavailable
     throw NotImplementedError
  for tuple in fwalk(args):
     yield tuple[:-1]
  return

+
+def walk(top=".", topdown=True, onerror=None, followlinks=False, *, dir_fd=None):
     """Directory tree generator.
 
     For each directory in the directory tree rooted at top (including top
     itself, but excluding '.' and '..'), yields a 3-tuple
 
         dirpath, dirnames, filenames
 
     dirpath is a string, the path to the directory.  dirnames is a list of
     the names of the subdirectories in dirpath (excluding '.' and '..').
     filenames is a list of the names of the non-directory files in dirpath.
@@ skipping 43 matching lines @@
             dirs.remove('CVS')  # don't visit CVS directories
     """
 
     islink, join, isdir = path.islink, path.join, path.isdir
 
     # We may not have read permission for top, in which case we can't
     # get a list of the files the directory contains.  os.walk
     # always suppressed the exception then, rather than blow up for a
     # minor reason when (say) a thousand readable directories are still
     # left to visit.  That logic is copied here.
+    #
+    # Note that listdir and error are globals in this module due
+    # to earlier import-*.
     try:
-        # Note that listdir and error are globals in this module due
-        # to earlier import-*.
-        names = listdir(top)
+        if dir_fd is None:
+            close_me = None
+            names = listdir(top)
+        else:
+            close_me = open(top, O_RDONLY, dir_fd=dir_fd)
+            names = listdir(close_me)
     except error as err:
+        if close_me is not None:
+            close(close_me)

[storchaka, 2012/06/25 14:08:10]

What if error happens on close?

[larry, 2012/06/25 15:04:13]

Obviated by new implementation.

         if onerror is not None:
             onerror(err)
         return
 
     dirs, nondirs = [], []
-    for name in names:
-        if isdir(join(top, name)):
-            dirs.append(name)
-        else:
-            nondirs.append(name)
+    if close_me is not None:
+        for name in names:
+            # name is in directory referenced by close_me!
+            # don't bother joining with top here.
+            if _isdir_dir_fd(name, close_me):
+                dirs.append(name)
+            else:
+                nondirs.append(name)
+        close(close_me)
+    else:
+        for name in names:
+            new_path = join(top, name)
+            if isdir(new_path):
+                dirs.append(name)
+            else:
+                nondirs.append(name)
 
     if topdown:
         yield top, dirs, nondirs
     for name in dirs:
         new_path = join(top, name)
         if followlinks or not islink(new_path):

[storchaka, 2012/06/25 14:08:10]

st.S_ISLINK(stat(name, dir_fd=close_me).st_mode) ?

[larry, 2012/06/25 15:04:13]

Obviated by new implementation.

-            yield from walk(new_path, topdown, onerror, followlinks)
+            yield from walk(new_path, topdown, onerror, followlinks, dir_fd=dir_fd)
     if not topdown:
         yield top, dirs, nondirs
 
 __all__.append("walk")
+if (listdir in supports_fd) and ({open, stat} <= supports_dir_fd):
+    supports_dir_fd.add(walk)
 
-if open in supports_dir_fd:
+if {open, stat} <= supports_dir_fd and {listdir, stat} <= supports_fd:
 
-    def fwalk(top, topdown=True, onerror=None, followlinks=False):
+    def fwalk(top=".", topdown=True, onerror=None, followlinks=False, *, dir_fd=None):
         """Directory tree generator.
 
         This behaves exactly like walk(), except that it yields a 4-tuple
 
             dirpath, dirnames, filenames, dirfd
 
         `dirpath`, `dirnames` and `filenames` are identical to walk() output,
         and `dirfd` is a file descriptor referring to the directory `dirpath`.
 
-        The advantage of walkfd() over walk() is that it's safe against symlink
+        The advantage of fwalk() over walk() is that it's safe against symlink
         races (when followlinks is False).
+
+        If dir_fd is not None, it should be a file descriptor open to a directory,
+          and top should be relative; top will then be relative to that directory.
+          (dir_fd is always supported for fwalk.)
 
         Caution:
         Since fwalk() yields file descriptors, those are only valid until the
         next iteration step, so you should dup() them if you want to keep them
         for a longer period.
 
         Example:
 
         import os
         for root, dirs, files, rootfd in os.fwalk('python/Lib/email'):
             print(root, "consumes", end="")
             print(sum([os.stat(name, dir_fd=rootfd).st_size for name in files]),
                   end="")
             print("bytes in", len(files), "non-directory files")
             if 'CVS' in dirs:
                 dirs.remove('CVS')  # don't visit CVS directories
         """
         # Note: To guard against symlink races, we use the standard
         # lstat()/open()/fstat() trick.
-        orig_st = lstat(top)
-        topfd = open(top, O_RDONLY)
+        orig_st = stat(top, follow_symlinks=False, dir_fd=dir_fd)
+        topfd = open(top, O_RDONLY, dir_fd=dir_fd)
         try:
             if (followlinks or (st.S_ISDIR(orig_st.st_mode) and
-                                path.samestat(orig_st, fstat(topfd)))):
+                                path.samestat(orig_st, stat(topfd)))):
                 yield from _fwalk(topfd, top, topdown, onerror, followlinks)
         finally:
             close(topfd)
 
     def _fwalk(topfd, toppath, topdown, onerror, followlinks):
         # 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)
@@ skipping 22 matching lines @@
 
         for name in dirs:
             try:
                 orig_st = stat(name, dir_fd=topfd, follow_symlinks=followlinks)
                 dirfd = open(name, O_RDONLY, dir_fd=topfd)
             except error as err:
                 if onerror is not None:
                     onerror(err)
                 return
             try:
-                if followlinks or path.samestat(orig_st, fstat(dirfd)):
+                if followlinks or path.samestat(orig_st, stat(dirfd)):
                     dirpath = path.join(toppath, name)
                     yield from _fwalk(dirfd, dirpath, topdown, onerror, followlinks)
             finally:
                 close(dirfd)
 
         if not topdown:
             yield toppath, dirs, nondirs, topfd
 
     __all__.append("fwalk")
 
@@ skipping 503 matching lines @@
         return getattr(self._stream, name)
     def __iter__(self):
         return iter(self._stream)
 
 # Supply os.fdopen()
 def fdopen(fd, *args, **kwargs):
     if not isinstance(fd, int):
         raise TypeError("invalid fd type (%s, expected integer)" % type(fd))
     import io
     return io.open(fd, *args, **kwargs)