contexts.py 2.1 KB

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