feat(util.dataset): add simulate_if_needed one-line auto-simulate helper

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
 
 

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"