contexts.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. import tempfile
  2. import os
  3. import shutil
  4. import sys
  5. import contextlib
  6. import site
  7. from ..compat import StringIO
  8. @contextlib.contextmanager
  9. def tempdir(cd=lambda dir:None, **kwargs):
  10. temp_dir = tempfile.mkdtemp(**kwargs)
  11. orig_dir = os.getcwd()
  12. try:
  13. cd(temp_dir)
  14. yield temp_dir
  15. finally:
  16. cd(orig_dir)
  17. shutil.rmtree(temp_dir)
  18. @contextlib.contextmanager
  19. def environment(**replacements):
  20. """
  21. In a context, patch the environment with replacements. Pass None values
  22. to clear the values.
  23. """
  24. saved = dict(
  25. (key, os.environ[key])
  26. for key in replacements
  27. if key in os.environ
  28. )
  29. # remove values that are null
  30. remove = (key for (key, value) in replacements.items() if value is None)
  31. for key in list(remove):
  32. os.environ.pop(key, None)
  33. replacements.pop(key)
  34. os.environ.update(replacements)
  35. try:
  36. yield saved
  37. finally:
  38. for key in replacements:
  39. os.environ.pop(key, None)
  40. os.environ.update(saved)
  41. @contextlib.contextmanager
  42. def quiet():
  43. """
  44. Redirect stdout/stderr to StringIO objects to prevent console output from
  45. distutils commands.
  46. """
  47. old_stdout = sys.stdout
  48. old_stderr = sys.stderr
  49. new_stdout = sys.stdout = StringIO()
  50. new_stderr = sys.stderr = StringIO()
  51. try:
  52. yield new_stdout, new_stderr
  53. finally:
  54. new_stdout.seek(0)
  55. new_stderr.seek(0)
  56. sys.stdout = old_stdout
  57. sys.stderr = old_stderr
  58. @contextlib.contextmanager
  59. def save_user_site_setting():
  60. saved = site.ENABLE_USER_SITE
  61. try:
  62. yield saved
  63. finally:
  64. site.ENABLE_USER_SITE = saved
  65. @contextlib.contextmanager
  66. def suppress_exceptions(*excs):
  67. try:
  68. yield
  69. except excs:
  70. pass