setenviron contextmanager source code

setenviron.py

import contextlib
import os


@contextlib.contextmanager
def setenviron(envdict=None, **mapping):
    """``with`` context to temporarily modify the environment variables"""
    _environ = os.environ.copy()

    if envdict:
        for key, value in envdict.items():
            os.environ[key] = str(value)

    for key, value in mapping.items():
        os.environ[key] = str(value)

    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(_environ)