A Python context manager for setting/unsetting environment variables

env-context.py

from contextlib import contextmanager

"""
Usage:

with env_var('MY_VAR', 'foo'):
    # is set here
    
# does not exist here
"""

@contextmanager
def env_var(key, value):
    os.environ[key] = value
    yield
    del os.environ[key]

This further improves over @jdemaeyer's snippet. It allows using the context manager to temporarily unset an env variable (e.g. with environ(HOME=None)).

Also, because it uses pop(), it will not choke if an environment variable that was added by the context manager has been deleted by the wrapped code.

import os
from contextlib import contextmanager


@contextmanager
def environ(**env):
    originals = {k: os.environ.get(k) for k in env}
    for k, val in env.items():
        if val is None:
            os.environ.pop(k, None)
        else:
            os.environ[k] = val
    try:
        yield
    finally:
        for k, val in originals.items():
            if val is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = val