def env_variables(**kwargs):
    """CM for temporarily setting (or removing) environment variables."""
    old_values = {name: os.environ.get(name, None) for name in kwargs}

    try:
        for name, new_value in kwargs.items():
            if new_value is not None:
                os.environ[name] = str(new_value)
            elif name in os.environ:
                os.environ.pop(name)
        yield
    finally:
        for name, old_value in old_values.items():
            if old_value is not None:
                os.environ[name] = old_value
            elif name in os.environ:
                os.environ.pop(name)
CM for temporarily setting (or removing) environment variables.
python/aitemplate/testing/test_utils.py