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