Python - 临时修改当前进程的环境

2022-01-15 00:00:00 python environment-variables

问题描述

我使用以下代码临时修改环境变量.

I use the following code to temporarily modify environment variables. 

@contextmanager
def _setenv(**mapping):
    """``with`` context to temporarily modify the environment variables"""
    backup_values = {}
    backup_remove = set()
    for key, value in mapping.items():
        if key in os.environ:
            backup_values[key] = os.environ[key]
        else:
            backup_remove.add(key)
        os.environ[key] = value

    try:
        yield
    finally:
        # restore old environment
        for k, v in backup_values.items():
            os.environ[k] = v
        for k in backup_remove:
            del os.environ[k]

这个with上下文主要用在测试用例中.例如，

This with context is mainly used in test cases. For example,

def test_myapp_respects_this_envvar():
    with _setenv(MYAPP_PLUGINS_DIR='testsandbox/plugins'):
        myapp.plugins.register()
        [...]

我的问题:有没有一种简单/优雅的方式来编写 _setenv?我考虑过实际执行 backup = os.environ.copy() 然后 os.environ = backup .. 但我不确定这是否会影响程序行为(例如:如果 os.environ 在 Python 解释器的其他地方被引用).

My question: is there a simple/elegant way to write _setenv? I thought about actually doing backup = os.environ.copy() and then os.environ = backup .. but I am not sure if that would affect the program behavior (eg: if os.environ is referenced elsewhere in the Python interpreter).

解决方案

_environ = dict(os.environ)  # or os.environ.copy()
try:

    ...

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

相关文章