1  # Copyright (c) Microsoft Corporation. All rights reserved.
   2  # Licensed under the MIT License. See LICENSE in the project root
   3  # for license information.
   4  
   5  import contextlib
   6  import os
   7  
   8  
   9  @contextlib.contextmanager
  10  def cwd(dirname):
  11      """A context manager for operating in a different directory."""
  12      orig = os.getcwd()
  13      os.chdir(dirname)
  14      try:
  15          yield orig
  16      finally:
  17          os.chdir(orig)
  18  
  19  
  20  def iter_all_files(root, prune_dir=None, exclude_file=None):
  21      """Yield (dirname, basename, filename) for each file in the tree.
  22  
  23      This is an alternative to os.walk() that flattens out the tree and
  24      with filtering.
  25      """
  26      pending = [root]
  27      while pending:
  28          dirname = pending.pop(0)
  29          for result in _iter_files(dirname, pending, prune_dir, exclude_file):
  30              yield result
  31  
  32  
  33  def iter_tree(root, prune_dir=None, exclude_file=None):
  34      """Yield (dirname, files) for each directory in the tree.
  35  
  36      The list of files is actually a list of (basename, filename).
  37  
  38      This is an alternative to os.walk() with filtering."""
  39      pending = [root]
  40      while pending:
  41          dirname = pending.pop(0)
  42          files = []
  43          for _, b, f in _iter_files(dirname, pending, prune_dir, exclude_file):
  44              files.append((b, f))
  45          yield dirname, files
  46  
  47  
  48  def _iter_files(dirname, subdirs, prune_dir, exclude_file):
  49      for basename in os.listdir(dirname):
  50          filename = os.path.join(dirname, basename)
  51          if os.path.isdir(filename):
  52              if prune_dir is not None and prune_dir(dirname, basename):
  53                  continue
  54              subdirs.append(filename)
  55          else:
  56              # TODO: Use os.path.isfile() to narrow it down?
  57              if exclude_file is not None and exclude_file(dirname, basename):
  58                  continue
  59              yield dirname, basename, filename
  60