Python - tạm thời sửa đổi môi trường của quy trình hiện tại

Question

Tôi sử dụng mã sau để tạm thời sửa đổi các biến môi trường.

@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] 

Ngữ cảnh with này chủ yếu được sử dụng trong các trường hợp thử nghiệm. Ví dụ:

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

Câu hỏi của tôi: có cách đơn giản/thanh lịch để viết _setenv? Tôi nghĩ về thực sự làm backup = os.environ.copy() và sau đó os.environ = backup .. nhưng tôi không chắc liệu điều đó có ảnh hưởng đến hành vi của chương trình hay không (ví dụ: nếu os.environ là được tham chiếu ở nơi khác trong trình thông dịch Python).

Answer 1

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

    ... 

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

Answer 2

tôi đề nghị bạn thực hiện như sau:

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: thực hiện tiến bộ hơn

Người quản lý bối cảnh dưới đây có thể được sử dụng để thêm/xóa/cập nhật biến môi trường của bạn:

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] 

Cách sử dụng ví dụ:

>>> 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

Để kiểm tra đơn vị Tôi thích sử dụng chức năng trang trí với thông số tùy chọn. Bằng cách này tôi có thể sử dụng các giá trị môi trường đã sửa đổi cho toàn bộ chức năng thử nghiệm. Các trang trí dưới đây cũng khôi phục các giá trị môi trường ban đầu trong trường hợp chức năng đặt ra một ngoại lệ:

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 

Cách sử dụng trong các thử nghiệm đơn vị:

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

Sử dụng các ý chính ở đây, bạn có thể lưu/khôi phục địa phương, phạm vi toàn cầu biến và môi trường biến: 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')) 

đầu ra này:

93 
94 
3 
4 
bar 
None