diff --git a/devito/parameters.py b/devito/parameters.py index 72cd3bb59d..5c60f0348e 100644 --- a/devito/parameters.py +++ b/devito/parameters.py @@ -295,10 +295,14 @@ class switchenv(SwitchDecorator): the context manager, so should be used cautiously. """ def __init__(self, params): - self.previous = dict(os.environ) self.params = params + self.previous = {} def __enter__(self): + # Snapshot the environment upon entering, not upon construction, since the + # same object is reused across entries, most notably as a decorator + self.previous = dict(os.environ) + # Prevent having multiple conflicting device vars, e.g # switching CUDA_VISIBLE_DEVICES but having NVIDIA_VISIBLE_DEVICES set. from devito.arch.archinfo import device_vars diff --git a/tests/test_tools.py b/tests/test_tools.py index 2996844ee5..37e4ffca81 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -344,3 +344,33 @@ def test_switchenv(): # Make sure the switchenv does not persist to verify switchenv works as intended assert dict(os.environ) == previous_environ + + +def test_switchenv_reuse(): + # Save previous environment + previous_environ = dict(os.environ) + + try: + # A switchenv is constructed once, when the decorator is applied, and then + # reused on every call of the decorated function + @switchenv({'TEST_VAR': 'foo'}) + def foo(): + return os.environ['TEST_VAR'] + + # Set after the decorator has been applied, so it is not visible to an + # environment snapshot taken at construction time + os.environ['TEST_VAR_LATE'] = 'bar' + + assert foo() == 'foo' + assert os.environ.get('TEST_VAR') is None + assert os.environ['TEST_VAR_LATE'] == 'bar' + + # Same story for a switchenv reused as a context manager + cm = switchenv({'TEST_VAR': 'foo'}) + os.environ['TEST_VAR_LATER'] = 'baz' + with cm: + assert os.environ['TEST_VAR'] == 'foo' + assert os.environ['TEST_VAR_LATER'] == 'baz' + finally: + os.environ.clear() + os.environ.update(previous_environ)