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

时间:2022-06-29 10:51:43

I use the following code to temporarily modify environment variables.

我使用以下代码临时修改环境变量。

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

This with context is mainly used in test cases. For example,

这与上下文主要用于测试用例。例如,

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

My question: is there a simple/elegant way to write _setenv? I thought about actually doing backup = os.environ.copy() and then os.environ = backup .. but I am not sure if that would affect the program behavior (eg: if os.environ is referenced elsewhere in the Python interpreter).

我的问题:是否有一种简单/优雅的方式来编写_setenv?我考虑过实际做backup = os.environ.copy()然后os.environ = backup ..但我不确定这是否会影响程序行为(例如:如果os.environ在Python解释器的其他地方被引用) 。

4 个解决方案

#1


23  

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

    ...

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

#2


25  

I suggest you the following implementation:

我建议你实施以下内容:

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: more advanced implementation

编辑:更高级的实施

The context manager below can be used to add/remove/update your environment variables:

下面的上下文管理器可用于添加/删除/更新您的环境变量:

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]

Usage examples:

用法示例:

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

#3


1  

For unit testing I prefer using a decorator function with optional parameters. This way I can use the modified environment values for a whole test function. The decorator below also restores the original environment values in case the function raises an Exception:

对于单元测试,我更喜欢使用带有可选参数的装饰器功能。这样我就可以将修改后的环境值用于整个测试功能。下面的装饰器还会恢复原始环境值,以防函数引发异常:

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

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

#4


0  

Using the gist here, you can save/restore local, global scope variable and environment variables: https://gist.github.com/earonesty/ac0617a5672ae1a41be1eaf316dd63e4

使用此处的要点,您可以保存/恢复本地,全局范围变量和环境变量: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'))

This outputs:

这输出:

93
94
3
4
bar
None

#1


23  

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

    ...

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

#2


25  

I suggest you the following implementation:

我建议你实施以下内容:

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: more advanced implementation

编辑:更高级的实施

The context manager below can be used to add/remove/update your environment variables:

下面的上下文管理器可用于添加/删除/更新您的环境变量:

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]

Usage examples:

用法示例:

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

#3


1  

For unit testing I prefer using a decorator function with optional parameters. This way I can use the modified environment values for a whole test function. The decorator below also restores the original environment values in case the function raises an Exception:

对于单元测试,我更喜欢使用带有可选参数的装饰器功能。这样我就可以将修改后的环境值用于整个测试功能。下面的装饰器还会恢复原始环境值,以防函数引发异常:

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

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

#4


0  

Using the gist here, you can save/restore local, global scope variable and environment variables: https://gist.github.com/earonesty/ac0617a5672ae1a41be1eaf316dd63e4

使用此处的要点,您可以保存/恢复本地,全局范围变量和环境变量: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'))

This outputs:

这输出:

93
94
3
4
bar
None