2010-01-13 13 views
13

Ich verwende den folgenden Code, um vorübergehend Umgebungsvariablen zu ändern.Python - vorübergehend die Umgebung des aktuellen Prozesses ändern

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

Dieser with Kontext wird hauptsächlich in Testfällen verwendet. Zum Beispiel

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

Meine Frage: Gibt es eine einfache/elegante Weise _setenv zu schreiben? Ich dachte darüber nach, backup = os.environ.copy() und dann os.environ = backup zu tun .. aber ich bin mir nicht sicher, ob das das Programmverhalten beeinflussen würde (zB: wenn os.environ ist verwiesen an anderer Stelle im Python-Interpreter).

Antwort

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

    ... 

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

Gut. Ich benutze '.copy()' anstelle von 'dict()' obwohl. –

+0

wollte ich nur, danke! – nnachefski

+2

Ok, aber im Falle eines Fehlers (Ausnahme) während [...] werden Umgebungsvariablen nicht wiederhergestellt: ein 'try ... finally ...' ist dafür erforderlich. –

21

Ich schlage vor, Sie die folgende Implementierung:

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: erweiterte Implementierung

Der Kontext-Manager unten verwendet werden können, hinzufügen/entfernen/aktualisieren Sie Ihre Umgebungsvariablen:

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] 

Verwendung Beispiele:

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

Für Besucher dieser alten Frage, sehe ich keine offensichtlichen Mängel in dieser Antwort und es ist vollständiger und nützlicher als das Original. – KobeJohn

+0

Dies sollte Teil von Python sein - oder so etwas. Die Umgebung zu testen ist fies - aber manchmal notwendig - Zeug und kann die Tests, die hinter env-messenden Test-Funktionen liegen, ernsthaft unterbrechen, entkräften oder anderweitig stören :( – Chris

+0

Das ist die bessere Antwort :) –

0

Für Komponententests bevorzuge ich eine Decorator-Funktion mit optionalen Parametern. Auf diese Weise kann ich die modifizierten Umgebungswerte für eine ganze Testfunktion verwenden. Der Dekorateur unten auch wieder die ursprünglichen Umgebungswerte, falls die Funktion eine Ausnahme auslöst:

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 

Verwendung in Unit-Tests:

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() 
0

das Wesentliche hier verwenden, können Sie lokalen, globalen Bereich Speicher/Wiederherstellen variable und Umgebungsvariablen: 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')) 

Diese Ausgänge:

93 
94 
3 
4 
bar 
None 
Verwandte Themen