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 .. но я не уверен, повлияет ли это на поведение программы (например: if 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) 

EDIT: более продвинутую реализацию

менеджер контекста ниже, может быть использован для добавления/удаления/обновить переменные окружения:

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