Python - temporarily modify the current process's environment

Question

I use the following code to temporarily modify environment variables.

@contextmanager def _setenv(**mapping):     """``with`` context to temporarily modify the environment variables"""     backup_values = {}     backup_remove = set()     for key, value in mapping.items():         if key in os.environ:             backup_values[key] = os.environ[key]         else:             backup_remove.add(key)         os.environ[key] = value      try:         yield     finally:         # restore old environment         for k, v in backup_values.items():             os.environ[k] = v         for k in backup_remove:             del os.environ[k] 

This with context is mainly used in test cases. For example,

def test_myapp_respects_this_envvar():     with _setenv(MYAPP_PLUGINS_DIR='testsandbox/plugins'):         myapp.plugins.register()         [...] 

My question: is there a simple/elegant way to write _setenv? I thought about actually doing backup = os.environ.copy() and then os.environ = backup .. but I am not sure if that would affect the program behavior (eg: if os.environ is referenced elsewhere in the Python interpreter).

Answer 1

I suggest you the following implementation:

import contextlib import os   @contextlib.contextmanager def set_env(**environ):     """     Temporarily set the process environment variables.      >>> with set_env(PLUGINS_DIR=u'test/plugins'):     ...   "PLUGINS_DIR" in os.environ     True      >>> "PLUGINS_DIR" in os.environ     False      :type environ: dict[str, unicode]     :param environ: Environment variables to set     """     old_environ = dict(os.environ)     os.environ.update(environ)     try:         yield     finally:         os.environ.clear()         os.environ.update(old_environ) 

EDIT: more advanced implementation

The context manager below can be used to add/remove/update your environment variables:

import contextlib import os   @contextlib.contextmanager def modified_environ(*remove, **update):     """     Temporarily updates the ``os.environ`` dictionary in-place.      The ``os.environ`` dictionary is updated in-place so that the modification     is sure to work in all situations.      :param remove: Environment variables to remove.     :param update: Dictionary of environment variables and values to add/update.     """     env = os.environ     update = update or {}     remove = remove or []      # List of environment variables being updated or removed.     stomped = (set(update.keys()) | set(remove)) & set(env.keys())     # Environment variables and values to restore on exit.     update_after = {k: env[k] for k in stomped}     # Environment variables and values to remove on exit.     remove_after = frozenset(k for k in update if k not in env)      try:         env.update(update)         [env.pop(k, None) for k in remove]         yield     finally:         env.update(update_after)         [env.pop(k) for k in remove_after] 

Usage examples:

>>> with modified_environ('HOME', LD_LIBRARY_PATH='/my/path/to/lib'): ...     home = os.environ.get('HOME') ...     path = os.environ.get("LD_LIBRARY_PATH") >>> home is None True >>> path '/my/path/to/lib'  >>> home = os.environ.get('HOME') >>> path = os.environ.get("LD_LIBRARY_PATH") >>> home is None False >>> path is None True 

EDIT2

A demonstration of this context manager is available on GitHub.

Answer 2

_environ = dict(os.environ)  # or os.environ.copy() try:      ...  finally:     os.environ.clear()     os.environ.update(_environ)