A context manager for temporarily setting environment variables. Handles resetting all env vars to their previous values.

env_var_context.py

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()