You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

83 lines
2.3 KiB

4 years ago
  1. import contextlib
  2. import os
  3. @contextlib.contextmanager
  4. def temporary_env(newenv):
  5. """Completely replace the environment variables with the specified dict.
  6. Use as a context manager::
  7. with temporary_env({'PATH': my_path}):
  8. ...
  9. """
  10. orig_env = os.environ.copy()
  11. os.environ.clear()
  12. os.environ.update(newenv)
  13. try:
  14. yield
  15. finally:
  16. os.environ.clear()
  17. os.environ.update(orig_env)
  18. @contextlib.contextmanager
  19. def modified_env(changes, snapshot=True):
  20. """Temporarily modify environment variables.
  21. Specify the changes as a dictionary mapping names to new values, using
  22. None as the value for names that should be deleted.
  23. Example use::
  24. with modified_env({'SHELL': 'bash', 'PYTHONPATH': None}):
  25. ...
  26. When the context exits, there are two possible ways to restore the
  27. environment. If *snapshot* is True, the default, it will reset the whole
  28. environment to its state when the context was entered. If *snapshot* is
  29. False, it will restore only the specific variables it modified, leaving
  30. any changes made to other environment variables in the context.
  31. """
  32. def update_del(changes):
  33. for k, v in changes.items():
  34. if v is None:
  35. os.environ.pop(k, None)
  36. else:
  37. os.environ[k] = v
  38. if snapshot:
  39. saved_variables = os.environ.copy()
  40. else:
  41. saved_variables = {}
  42. for k,v in changes.items():
  43. saved_variables[k] = os.environ.get(k, None)
  44. update_del(changes)
  45. try:
  46. yield
  47. finally:
  48. if snapshot:
  49. os.environ.clear()
  50. os.environ.update(saved_variables)
  51. else:
  52. update_del(saved_variables)
  53. def make_env_restorer():
  54. """Snapshot the current environment, return a function to restore that.
  55. This is intended to produce cleanup functions for tests. For example,
  56. using the :class:`unittest.TestCase` API::
  57. def setUp(self):
  58. self.addCleanup(testpath.make_env_restorer())
  59. Any changes a test makes to the environment variables will be wiped out
  60. before the next test is run.
  61. """
  62. orig_env = os.environ.copy()
  63. def restore():
  64. os.environ.clear()
  65. os.environ.update(orig_env)
  66. return restore