A Python context manager for setting/unsetting environment variables

env-context.py

"""
Environment variable context manager
------------------------------------

Support utility for managing environments in which e.g. git is run

"""
import os
from contextlib import contextmanager
from functools import wraps

@contextmanager
def env_var(key, value):
    """
    Set environment variable for the context

    Example:
    --------

        with env_var('GIT_SSL_NO_VERIFY', '1'):
            # Environment variable is set here
            git.Repo.clone_from('https://giturl/user/project.git', 'some/dir')

        # Environment variable is reset here

    :param key: name of environment variable in dict :var:``os.environ``
    :param value: value to set for the context
    """
    old_value = os.environ.get(key, None)

    os.environ[key] = value

    yield

    if old_value is None:
        del os.environ[key]
    else:
        os.environ[key] = old_value


def with_env_var(key, value):
    """
    Returns a decorator which performs the decorated function with the environment variable set.

    The environment variable is reset when the function returns

    :param key: name of environment variable in dict :var:``os.environ``
    :param value: value to set for the context
    :return: decorator
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            with env_var(key, value):
                return func(*args, **kwargs)
        return wrapper

    return decorator