_util.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  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. import contextlib
  5. import os
  6. @contextlib.contextmanager
  7. def cwd(dirname):
  8. """A context manager for operating in a different directory."""
  9. orig = os.getcwd()
  10. os.chdir(dirname)
  11. try:
  12. yield orig
  13. finally:
  14. os.chdir(orig)
  15. def iter_all_files(root, prune_dir=None, exclude_file=None):
  16. """Yield (dirname, basename, filename) for each file in the tree.
  17. This is an alternative to os.walk() that flattens out the tree and
  18. with filtering.
  19. """
  20. pending = [root]
  21. while pending:
  22. dirname = pending.pop(0)
  23. for result in _iter_files(dirname, pending, prune_dir, exclude_file):
  24. yield result
  25. def iter_tree(root, prune_dir=None, exclude_file=None):
  26. """Yield (dirname, files) for each directory in the tree.
  27. The list of files is actually a list of (basename, filename).
  28. This is an alternative to os.walk() with filtering."""
  29. pending = [root]
  30. while pending:
  31. dirname = pending.pop(0)
  32. files = []
  33. for _, b, f in _iter_files(dirname, pending, prune_dir, exclude_file):
  34. files.append((b, f))
  35. yield dirname, files
  36. def _iter_files(dirname, subdirs, prune_dir, exclude_file):
  37. for basename in os.listdir(dirname):
  38. filename = os.path.join(dirname, basename)
  39. if os.path.isdir(filename):
  40. if prune_dir is not None and prune_dir(dirname, basename):
  41. continue
  42. subdirs.append(filename)
  43. else:
  44. # TODO: Use os.path.isfile() to narrow it down?
  45. if exclude_file is not None and exclude_file(dirname, basename):
  46. continue
  47. yield dirname, basename, filename