misc: Snapshot the environment when a switchenv is entered

`switchenv.__init__` took the environment snapshot used by `__exit__` to
restore the environment, so the snapshot dated from construction rather than
from entry (devito/parameters.py, `switchenv.__init__`). `SwitchDecorator.__call__`
wraps the decorated function in `with self`, so a `switchenv` used as a
decorator is constructed once and entered on every call. Because `__exit__`
does `os.environ.clear()` followed by `os.environ.update(self.previous)`, the
first call reverted the environment to how it looked when the decorator was
applied, silently dropping every variable set since then. The same happened
for a `switchenv` object stored and reused as a context manager.

The sibling `switchconfig` already rebuilds `self.previous` inside `__enter__`.
This does the same for `switchenv`, taking the snapshot at the top of
`__enter__` before the device vars are popped, so the restored environment is
unchanged for a single use.

Author: MaxFreedomPollard

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)