Skip to content

misc: Snapshot the environment when a switchenv is entered - #3019

Open
MaxFreedomPollard wants to merge 1 commit into
devitocodes:mainfrom
MaxFreedomPollard:fix-switchenv-env-snapshot
Open

misc: Snapshot the environment when a switchenv is entered#3019
MaxFreedomPollard wants to merge 1 commit into
devitocodes:mainfrom
MaxFreedomPollard:fix-switchenv-env-snapshot

Conversation

@MaxFreedomPollard

Copy link
Copy Markdown

switchenv.__exit__ restores the environment with os.environ.clear() followed by os.environ.update(self.previous), but self.previous is captured in switchenv.__init__ (devito/parameters.py), so it dates from construction rather than from entry. SwitchDecorator.__call__ wraps the decorated function in with self, so a switchenv used as a decorator is built once, when the decorator is applied, and entered on every call. The first call therefore rewinds the environment to how it looked at decoration time, deleting every variable set since. Same story for a switchenv object kept around and reused as a context manager.

On main:

import os
from devito import switchenv

@switchenv({'FOO': 'inside'})
def f():
    return os.environ['FOO']

os.environ['SET_LATER'] = 'keepme'
f()
print(os.environ.get('SET_LATER'))  # None, was 'keepme'

The sibling switchconfig already rebuilds self.previous at the top of __enter__. This does the same in switchenv, taking the snapshot before the device vars are popped so that a single use restores exactly what it restored before. Every existing call site builds the object inline in a with statement, where construction and entry are adjacent, so none of them change behaviour.

Verified on macOS with DEVITO_ARCH=clang. The new tests/test_tools.py::test_switchenv_reuse fails on unmodified main with KeyError: 'TEST_VAR_LATE' and passes with the fix. pytest tests/test_tools.py gives 32 passed and pytest tests/test_environment.py gives 1 passed. isort --check-only ., ruff check --preview, flake8 --builtins=ArgumentError . and typos are all clean.

`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.
@codecov

codecov Bot commented Sep 7, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.38%. Comparing base (b330c14) to head (e5e28e4).

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3019      +/-   ##
==========================================
- Coverage   83.72%   79.38%   -4.35%     
==========================================
  Files         257      257              
  Lines       54822    54840      +18     
  Branches     4693     4693              
==========================================
- Hits        45901    43532    -2369     
- Misses       8110    10450    +2340     
- Partials      811      858      +47     
Flag Coverage Δ
pytest-gpu-aomp-amdgpuX ?
pytest-gpu-gcc- 77.22% <100.00%> (-1.12%) ⬇️
pytest-gpu-icx- ?
pytest-gpu-nvc-nvidiaX ?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@EdCaunt

EdCaunt commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

Needs a rebase, but then it's good to go

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants