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]

With some improvements:

• support multiple keys
• restore old values
• rollback imune to crash in yield

from contextlib import contextmanager
@contextmanager
def environ(**kwargs):
    orig = dict()
    todel = []
    try:
        for k, newval in kwargs.items():
            if k in os.environ:
                orig[k] = os.environ[k]
            else:
                todel.append(k)
            os.environ[k] = newval
        yield
    finally:
        for k, oldval in orig.items():
            os.environ[k] = oldval
        for k in todel:
            del os.environ[k]

# Usage
import unittest
class TestEnvMgr(unittest.TestCase):
    def test_with_environ(self):
       with environ(USER='xxx', HOME='42', NONEXISTENT='?'):
            self.assertEqual(os.environ['USER'], 'xxx')
            self.assertEqual(os.environ['HOME'], '42')
            self.assertEqual(os.environ['NONEXISTENT'], '?')
        self.assertEqual(os.environ['USER'], user)
        self.assertEqual(os.environ['HOME'], home)
        self.assertTrue('NONEXISTENT' not in os.environ)

Building on @ArnaudPel's version above, a little more concise (and using that existing environment vars can never be None, only ''):

from contextlib import contextmanager


@contextmanager
def environ(env):
    """Temporarily set environment variables inside the context manager and
    fully restore previous environment afterwards
    """
    original_env = {key: os.getenv(key) for key in env}
    os.environ.update(env)
    try:
        yield
    finally:
        for key, value in original_env.items():
            if value is None:
                del os.environ[key]
            else:
                os.environ[key] = value

Whether to use env or **env in the signature is a bit up to your taste, I prefer passing a dictionary rather than keyword args :)

This looks great, and seems like exactly what I need. But why not use the built-in patch.dict?
https://docs.python.org/3/library/unittest.mock.html#unittest.mock.patch.dict

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