|
|
- import contextlib
- import os
-
- @contextlib.contextmanager
- def temporary_env(newenv):
- """Completely replace the environment variables with the specified dict.
-
- Use as a context manager::
-
- with temporary_env({'PATH': my_path}):
- ...
- """
- orig_env = os.environ.copy()
- os.environ.clear()
- os.environ.update(newenv)
- try:
- yield
- finally:
- os.environ.clear()
- os.environ.update(orig_env)
-
- @contextlib.contextmanager
- def modified_env(changes, snapshot=True):
- """Temporarily modify environment variables.
-
- Specify the changes as a dictionary mapping names to new values, using
- None as the value for names that should be deleted.
-
- Example use::
-
- with modified_env({'SHELL': 'bash', 'PYTHONPATH': None}):
- ...
-
- When the context exits, there are two possible ways to restore the
- environment. If *snapshot* is True, the default, it will reset the whole
- environment to its state when the context was entered. If *snapshot* is
- False, it will restore only the specific variables it modified, leaving
- any changes made to other environment variables in the context.
- """
- def update_del(changes):
- for k, v in changes.items():
- if v is None:
- os.environ.pop(k, None)
- else:
- os.environ[k] = v
-
- if snapshot:
- saved_variables = os.environ.copy()
- else:
- saved_variables = {}
- for k,v in changes.items():
- saved_variables[k] = os.environ.get(k, None)
-
- update_del(changes)
-
- try:
- yield
- finally:
- if snapshot:
- os.environ.clear()
- os.environ.update(saved_variables)
- else:
- update_del(saved_variables)
-
- def make_env_restorer():
- """Snapshot the current environment, return a function to restore that.
-
- This is intended to produce cleanup functions for tests. For example,
- using the :class:`unittest.TestCase` API::
-
- def setUp(self):
- self.addCleanup(testpath.make_env_restorer())
-
- Any changes a test makes to the environment variables will be wiped out
- before the next test is run.
- """
- orig_env = os.environ.copy()
-
- def restore():
- os.environ.clear()
- os.environ.update(orig_env)
-
- return restore
|