From 94b9149d421f4cd0277ec305dd72ec836cc0a404 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Nov 2026 15:37:45 +0000 Subject: [PATCH] feat(util.dataset): add simulate_if_needed one-line auto-simulate helper Adds aa.util.dataset.simulate_if_needed(dataset_path, simulator), which runs the simulator script in a subprocess (sys.executable, check=True) only when should_simulate(dataset_path) is True. Subprocess isolation is kept deliberately (not runpy) so simulator JAX/matplotlib state cannot leak into the caller. Accepts str or Path for both args and returns whether the simulator ran. The should_simulate docstring's usage example now points at the helper. Tests mock subprocess.run for both branches and both arg types. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_016nhQ8NyQo2s3YUk8eUrrB3 --- autoarray/util/dataset_util.py | 58 ++++++++++++++++++++++-- test_autoarray/util/test_dataset_util.py | 55 ++++++++++++++++++++++ 2 files changed, 109 insertions(+), 4 deletions(-) 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"