Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
A context manager for temporarily setting environment variables. Handles resetting all env vars to their previous values.
import os
from typing import ContextManager, Dict, Optional, Union
class environment(ContextManager):
"""Context manager for temporarily setting environment variables. Resets env var state on completion.
**WARNING**: Mutates the `os.environ` environment variable mapping in the :func:`__enter__`
and :func:`__exit__` methods. Every `__enter__` **MUST** be followed-up by an `__exit__`.
NOTE: The suggested use pattern is to make a new :class:`environment` for each time one
requires the functionality. It is possible, however, to create one :class:`environment`
instance and re-use to temporarily set the same environment variable state.
For example, this is the most common use case:
>>>> with environment(ENV_VAR='your-value'):
>>>> # do something that needs this ENV_VAR set to 'your-value'
>>>> ...
>>>> # ENV_VAR is reset to prior value
"""
def __init__(self, **env_vars) -> None:
"""Keep track of the temporary values one will assign to a set of environment variables.
NOTE: Either supply string-valued keyword arguments or a dictionary of env var names to their values.
Environment variables must be either `str` or `int` valued.
Raises :class:`ValueError` if either of these conditions are invalid.
"""
# used in __enter__ and __exit__: keeps track of prior existing environment
# variable settings so they can be reset when the context block is finished
self.__prior_env_var_setting: Dict[str, Optional[Union[str, int]]] = dict()
# we store the user's set of desired temporary values for environment variables
# we also validate these settings to a minimum bar of correctness
self.__new_env_var_settings: Dict[str, Union[int, str]] = dict()
for env, val in env_vars.items():
# validate input
if not isinstance(env, str) or len(env) == 0:
raise ValueError(f"Need valid env var name, not ({type(env)}) '{env}'")
if not isinstance(val, (str, int)):
raise ValueError(f"Need valid env var value, not ({type(val)}) '{val}'")
self.__new_env_var_settings[env] = val
def __enter__(self) -> "environment":
"""Temporarily sets user's supplied environment variables.
Records the previous values of all environment variables to be set.
WARNING: Mutates internal sate.
"""
# get the existing values for all environment variables to be temporarily set
# set the env vars to their desired values too
for env, val in self.__new_env_var_settings.items():
# track prior value and set new for environment variable
prior: Optional[str] = os.environ.get(env)
self.__prior_env_var_setting[env] = prior
os.environ[env] = val
return self
def __exit__(self, exc_type, exc_value, traceback) -> None:
"""Restores the previous values of all temporarily set environment variables.
WARNING: Mutates internal sate.
NOTE: Ignores all input arguments.
"""
# restore all env vars
for env, prior in self.__prior_env_var_setting.items():
if prior is not None:
# restore previous environment variable value
os.environ[env] = prior
else:
# don't set to None if there wasn't a valid env var setting beforehand
del os.environ[env]
# forget about previous context setting
# could be used for another __enter__ --> __exit__ cycle
self.__prior_env_var_setting.clear()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment