diff --git a/autoarray/util/dataset_util.py b/autoarray/util/dataset_util.py index 7eb3a145c..0678b4c32 100644 --- a/autoarray/util/dataset_util.py +++ b/autoarray/util/dataset_util.py @@ -535,11 +535,16 @@ def should_simulate(dataset_path): then loaded silently, producing deterministic, environment-only failures that could not be reproduced in CI (autolens_workspace_test#260). - Use this as a drop-in replacement for ``not path.exists(dataset_path)`` in - the workspace auto-simulation pattern:: + Workspace scripts should not call this directly to auto-simulate their + dataset; use :func:`simulate_if_needed`, which wraps this predicate and the + simulator subprocess in one line:: - if aa.util.dataset.should_simulate(dataset_path): - subprocess.run([sys.executable, "scripts/.../simulator.py"], check=True) + aa.util.dataset.simulate_if_needed( + dataset_path, simulator="scripts/.../simulator.py" + ) + + Call this predicate directly only when the "simulate" step is something + other than running a simulator script. Precedence ---------- @@ -663,6 +668,51 @@ def should_simulate(dataset_path): return not Path(dataset_path).exists() + +def simulate_if_needed(dataset_path, simulator): + """ + Run the ``simulator`` script if the dataset at ``dataset_path`` needs to be + (re)simulated, as decided by :func:`should_simulate`. + + This is the one-line form of the workspace auto-simulation pattern:: + + aa.util.dataset.simulate_if_needed( + dataset_path, simulator="scripts/simulator/lens_x2.py" + ) + + The simulator is run in a **separate Python process** (``sys.executable``), + not in-process via ``runpy``: a simulator configures JAX, matplotlib and + module-level state that must not leak into the calling script. The run is + ``check=True``, so a failing simulator raises + ``subprocess.CalledProcessError`` rather than letting the caller load a + missing dataset. + + Parameters + ---------- + dataset_path + The dataset directory (``str`` or ``Path``) passed to + :func:`should_simulate` -- which may delete it, see that function. + simulator + Path (``str`` or ``Path``) to the simulator script, resolved relative to + the current working directory exactly as the ``subprocess.run`` call it + replaces was. + + Returns + ------- + bool + ``True`` if the simulator was run, ``False`` if the dataset was kept. + """ + if not should_simulate(str(dataset_path)): + return False + + import subprocess + import sys + + subprocess.run([sys.executable, str(simulator)], check=True) + + return True + + SMALL_DATASETS_N_CATALOGUE = 25 diff --git a/test_autoarray/util/test_dataset_util.py b/test_autoarray/util/test_dataset_util.py index 3ee2023ae..4d263e171 100644 --- a/test_autoarray/util/test_dataset_util.py +++ b/test_autoarray/util/test_dataset_util.py @@ -857,3 +857,58 @@ def test__cap_card__written_by_the_stack_is_read_back(monkeypatch, tmp_path): SMALL_DATASETS_SHAPE_NATIVE ) assert should_simulate(str(dataset_path)) is False + + +@pytest.mark.parametrize("as_path", [False, True]) +def test__simulate_if_needed__runs_simulator_only_when_should_simulate( + monkeypatch, tmp_path, as_path +): + import subprocess + import sys + + from autoarray.util import dataset_util + + calls = [] + monkeypatch.setattr(subprocess, "run", lambda *a, **kw: calls.append((a, kw))) + + dataset_path = tmp_path / "dataset" + simulator = tmp_path / "scripts" / "simulator.py" + if not as_path: + dataset_path, simulator = str(dataset_path), str(simulator) + + seen = [] + + def fake_should_simulate(path): + seen.append(path) + return True + + monkeypatch.setattr(dataset_util, "should_simulate", fake_should_simulate) + + assert aa.util.dataset.simulate_if_needed(dataset_path, simulator) is True + assert seen == [str(dataset_path)] + assert calls == [(([sys.executable, str(simulator)],), {"check": True})] + + monkeypatch.setattr(dataset_util, "should_simulate", lambda path: False) + calls.clear() + + assert ( + aa.util.dataset.simulate_if_needed(dataset_path, simulator=simulator) + is False + ) + assert calls == [] + + +def test__simulate_if_needed__missing_dataset__runs_simulator( + monkeypatch, tmp_path +): + import subprocess + + monkeypatch.delenv("PYAUTO_SMALL_DATASETS", raising=False) + + calls = [] + monkeypatch.setattr(subprocess, "run", lambda *a, **kw: calls.append(a[0])) + + assert aa.util.dataset.simulate_if_needed( + tmp_path / "missing", simulator="scripts/simulator.py" + ) + assert calls[0][1] == "scripts/simulator.py"