临时修改当前进程的环境

Question

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

@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 上下文主要用于测试案例。例如，

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 解释器的其他地方被 引用）。

Answer 1

我想做的事情和你一样，不过是为了单元测试。下面是我使用unittest.mock.patch函数实现的方法：

def test_function_with_different_env_variable():
    with mock.patch.dict('os.environ', {'hello': 'world'}, clear=True):
        self.assertEqual(os.environ.get('hello'), 'world')
        self.assertEqual(len(os.environ), 1)

简单来说，使用unittest.mock.patch.dict并设置clear=True，我们就把os.environ变成了一个只包含{'hello': 'world'}的字典。

• 如果去掉clear=True，那么原来的os.environ会保留，并且会在里面添加或替换指定的键值对{'hello': 'world'}。

• 如果去掉{'hello': 'world'}，那么就会创建一个空字典，这样在with块内，os.environ就会变成空的。

Answer 2

我建议你使用以下实现方式：

import contextlib
import os


@contextlib.contextmanager
def set_env(**environ):
    """
    Temporarily set the process environment variables.

    >>> with set_env(PLUGINS_DIR='test/plugins'):
    ...   "PLUGINS_DIR" in os.environ
    True

    >>> "PLUGINS_DIR" in os.environ
    False

    :type environ: dict[str, unicode]
    :param environ: Environment variables to set
    """
    old_environ = dict(os.environ)
    os.environ.update(environ)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(old_environ)

编辑：更高级的实现

下面的上下文管理器可以用来添加、删除或更新你的环境变量：

import contextlib
import os


@contextlib.contextmanager
def modified_environ(*remove, **update):
    """
    Temporarily updates the ``os.environ`` dictionary in-place.

    The ``os.environ`` dictionary is updated in-place so that the modification
    is sure to work in all situations.

    :param remove: Environment variables to remove.
    :param update: Dictionary of environment variables and values to add/update.
    """
    env = os.environ
    update = update or {}
    remove = remove or []

    # List of environment variables being updated or removed.
    stomped = (set(update.keys()) | set(remove)) & set(env.keys())
    # Environment variables and values to restore on exit.
    update_after = {k: env[k] for k in stomped}
    # Environment variables and values to remove on exit.
    remove_after = frozenset(k for k in update if k not in env)

    try:
        env.update(update)
        [env.pop(k, None) for k in remove]
        yield
    finally:
        env.update(update_after)
        [env.pop(k) for k in remove_after]

使用示例：

>>> with modified_environ('HOME', LD_LIBRARY_PATH='/my/path/to/lib'):
...     home = os.environ.get('HOME')
...     path = os.environ.get("LD_LIBRARY_PATH")
>>> home is None
True
>>> path
'/my/path/to/lib'

>>> home = os.environ.get('HOME')
>>> path = os.environ.get("LD_LIBRARY_PATH")
>>> home is None
False
>>> path is None
True

编辑2

这个上下文管理器的演示可以在 GitHub 上找到。

Answer 3

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

    ...

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

当然可以！请把你想要翻译的内容发给我，我会帮你用简单易懂的语言解释清楚。