关于Python：Python – 暂时修改当前进程的环境

Python - temporarily modify the current process's environment

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

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

1
2
3
4

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

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

我建议您执行以下操作：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

import contextlib
import os

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

>>> with set_env(PLUGINS_DIR=u'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)

编辑：更高级的实现

下面的上下文管理器可用于添加/删除/更新环境变量：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

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]

使用实例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14

>>> 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上提供了此上下文管理器的演示。