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

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

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

    ... 

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

Answer 2

我建议你实现如下：

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) 

编辑：更高级的实现

以下情况管理器可用于添加/删除/更新环境变量：

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 

Answer 3

对于单元测试，我更喜欢使用可选参数的装饰器函数。通过这种方式，我可以将修改后的环境值用于整个测试功能。下面也将装饰恢复的情况下，将抛出一个异常的原始环境值：

import os 

def patch_environ(new_environ=None, clear_orig=False): 
    if not new_environ: 
     new_environ = dict() 

    def actual_decorator(func): 
     from functools import wraps 

     @wraps(func) 
     def wrapper(*args, **kwargs): 
      original_env = dict(os.environ) 

      if clear_orig: 
       os.environ.clear() 

      os.environ.update(new_environ) 
      try: 
       result = func(*args, **kwargs) 
      except: 
       raise 
      finally: # restore even if Exception was raised 
       os.environ = original_env 

      return result 

     return wrapper 

    return actual_decorator 

使用单元测试：

class Something: 
    @staticmethod 
    def print_home(): 
     home = os.environ.get('HOME', 'unknown') 
     print("HOME = {0}".format(home)) 


class SomethingTest(unittest.TestCase): 
    @patch_environ({'HOME': '/tmp/test'}) 
    def test_environ_based_something(self): 
     Something.print_home() # prints: HOME = /tmp/test 

unittest.main() 

Answer 4

使用要点这里，你可以保存/恢复本地，全局范围变量和环境变量： https://gist.github.com/earonesty/ac0617a5672ae1a41be1eaf316dd63e4

import os 
from varlib import vartemp, envtemp 

x = 3 
y = 4 

with vartemp({'x':93,'y':94}): 
    print(x) 
    print(y) 
print(x) 
print(y) 

with envtemp({'foo':'bar'}): 
    print(os.getenv('foo')) 

print(os.getenv('foo')) 

此输出：

93 
94 
3 
4 
bar 
None