Merge pull request #3019 from MaxFreedomPollard/fix-switchenv-env-snapshot

misc: Snapshot the environment when a switchenv is entered

Author: EdCaunt

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

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)