Temporarily modify the current process's environment - environment-variables

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)

Answer 3

I was looking to do the same thing but for unit testing, here is how I have done it using the unittest.mock.patch function:
def test_function_with_different_env_variable():
with mock.patch.dict('os.environ', {'hello': 'world'}, clear=True):
self.assertEqual(os.environ.get('hello'), 'world')
self.assertEqual(len(os.environ), 1)
Basically using unittest.mock.patch.dict with clear=True, we are making os.environ as a dictionary containing solely {'hello': 'world'}.
Removing the clear=True will let the original os.environ and add/replace the specified key/value pair inside {'hello': 'world'}.
Removing {'hello': 'world'} will just create an empty dictionary, os.envrion will thus be empty within the with.

Answer 4

In pytest you can temporarily set an environment variable using the monkeypatch fixture. See the docs for details. I've copied a snippet here for your convenience.
import os
import pytest
from typing import Any, NewType
# Alias for the ``type`` of monkeypatch fixture.
MonkeyPatchFixture = NewType("MonkeyPatchFixture", Any)
# This is the function we will test below to demonstrate the ``monkeypatch`` fixture.
def get_lowercase_env_var(env_var_name: str) -> str:
"""
Return the value of an environment variable. Variable value is made all lowercase.
:param env_var_name:
The name of the environment variable to return.
:return:
The value of the environment variable, with all letters in lowercase.
"""
env_variable_value = os.environ[env_var_name]
lowercase_env_variable = env_variable_value.lower()
return lowercase_env_variable
def test_get_lowercase_env_var(monkeypatch: MonkeyPatchFixture) -> None:
"""
Test that the function under test indeed returns the lowercase-ified
form of ENV_VAR_UNDER_TEST.
"""
name_of_env_var_under_test = "ENV_VAR_UNDER_TEST"
env_var_value_under_test = "EnvVarValue"
expected_result = "envvarvalue"
# KeyError because``ENV_VAR_UNDER_TEST`` was looked up in the os.environ dictionary before its value was set by ``monkeypatch``.
with pytest.raises(KeyError):
assert get_lowercase_env_var(name_of_env_var_under_test) == expected_result
# Temporarily set the environment variable's value.
monkeypatch.setenv(name_of_env_var_under_test, env_var_value_under_test)
assert get_lowercase_env_var(name_of_env_var_under_test) == expected_result
def test_get_lowercase_env_var_fails(monkeypatch: MonkeyPatchFixture) -> None:
"""
This demonstrates that ENV_VAR_UNDER_TEST is reset in every test function.
"""
env_var_name_under_test = "ENV_VAR_UNDER_TEST"
expected_result = "envvarvalue"
with pytest.raises(KeyError):
assert get_lowercase_env_var(env_var_name_under_test) == expected_result

Answer 5

For unit testing I prefer using a decorator function with optional parameters. This way I can use the modified environment values for a whole test function. The decorator below also restores the original environment values in case the function raises an Exception:
import os
def patch_environ(new_environ=None, clear_orig=False):
if not new_environ:
new_environ = dict()
def actual_decorator(func):
from functools import wraps
#wraps(func)
def wrapper(*args, **kwargs):
original_env = dict(os.environ)
if clear_orig:
os.environ.clear()
os.environ.update(new_environ)
try:
result = func(*args, **kwargs)
except:
raise
finally: # restore even if Exception was raised
os.environ = original_env
return result
return wrapper
return actual_decorator
Usage in unit tests:
class Something:
#staticmethod
def print_home():
home = os.environ.get('HOME', 'unknown')
print("HOME = {0}".format(home))
class SomethingTest(unittest.TestCase):
#patch_environ({'HOME': '/tmp/test'})
def test_environ_based_something(self):
Something.print_home() # prints: HOME = /tmp/test
unittest.main()