From 474b6c1188c578306a5203e35c8e5ad2067d4c18 Mon Sep 17 00:00:00 2001 From: bruno-f-cruz <7049351+bruno-f-cruz@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:46:28 -0700 Subject: [PATCH] Allow explicit resolution order for `@experiment` --- src/clabe/cli.py | 9 +++++- src/clabe/launcher/_experiments.py | 47 +++++++++++++++++++++++++----- tests/launcher/test_experiments.py | 40 +++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 9 deletions(-) diff --git a/src/clabe/cli.py b/src/clabe/cli.py index e399822e..f3b2a796 100644 --- a/src/clabe/cli.py +++ b/src/clabe/cli.py @@ -32,11 +32,18 @@ class _RunCli(LauncherCliArgs): default=False, description="Refuse to start if another CLABE session is already running (used when serving the web UI)", ) + experiment_name: str | None = Field( + default=None, + description="Name of the @experiment to run when the file defines more than one. " + "If omitted, prompts interactively (or runs the only one found).", + ) def _run(self): """Builds the launcher, selects the experiment and runs it.""" launcher = Launcher(settings=self) - experiment_metadata = _select_experiment(self.experiment_path, frontend=launcher.frontend) + experiment_metadata = _select_experiment( + self.experiment_path, frontend=launcher.frontend, experiment_name=self.experiment_name + ) launcher.run_experiment(experiment_metadata.func) def cli_cmd(self): diff --git a/src/clabe/launcher/_experiments.py b/src/clabe/launcher/_experiments.py index 73979c7c..04bb2f64 100644 --- a/src/clabe/launcher/_experiments.py +++ b/src/clabe/launcher/_experiments.py @@ -23,10 +23,13 @@ class ExperimentMetadata: Attributes: name: Human-readable name for the experiment. func: The underlying callable. + order: Sort key controlling position when multiple experiments are + listed (lower sorts first). Ties keep declaration order. """ name: str func: ExperimentCallable + order: int = 0 class _IExperiment(Protocol): @@ -40,12 +43,20 @@ def __call__(self, launcher: "Launcher", /, *args: Any, **kwargs: Any) -> Any: . def experiment( *, name: str | None = None, + order: int = 0, ) -> Callable[[_IExperiment], _IExperiment]: """Decorator to mark a function as a CLABE experiment. The decorated function must accept a single `Launcher` argument and may be either synchronous or asynchronous. + Args: + name: Human-readable name for the experiment. Defaults to the + function's ``__name__``. + order: Sort key controlling where this experiment appears when a + module defines more than one (lower sorts first). Experiments with + the same ``order`` keep their declaration order. + Example: ```python from pathlib import Path @@ -54,7 +65,7 @@ def experiment( from clabe.launcher import experiment - @experiment(name="super_duper_experiment") + @experiment(name="super_duper_experiment", order=-1) async def vr_foraging_with_photometry(launcher: Launcher) -> None: ... ``` @@ -65,6 +76,7 @@ def decorator(func: _IExperiment) -> _IExperiment: metadata = ExperimentMetadata( name=exp_name, func=func, # type: ignore[arg-type] + order=order, ) func.__clabe_experiment_metadata__ = metadata return func @@ -92,13 +104,19 @@ def get_experiment_name(experiment: _IExperiment) -> str | None: def collect_clabe_experiments(module: ModuleType) -> Iterable[ExperimentMetadata]: - """Yield all `@experiment` experiments defined in the target module.""" + """Yield all `@experiment` experiments defined in the target module. + Experiments are yielded sorted by their ``order`` (lower first); ties keep + the module's declaration order. + """ + + discovered: list[ExperimentMetadata] = [] for value in vars(module).values(): metadata = getattr(value, "__clabe_experiment_metadata__", None) if isinstance(metadata, ExperimentMetadata): logger.debug("Discovered CLABE experiment: %s in module %s", metadata.name, module.__name__) - yield metadata + discovered.append(metadata) + yield from sorted(discovered, key=lambda e: e.order) def _load_module_from_path(path: Path): @@ -125,13 +143,17 @@ def _load_module_from_path(path: Path): return module -def _select_experiment(file_path: Path, frontend: Frontend | None = None) -> ExperimentMetadata: +def _select_experiment( + file_path: Path, frontend: Frontend | None = None, experiment_name: str | None = None +) -> ExperimentMetadata: """Select an experiment callable from a Python module. Loads the module at ``file_path``, discovers all callables decorated with :func:`experiment`, and returns the associated :class:`ExperimentMetadata`. - If a single experiment is found it is returned directly. When multiple + If ``experiment_name`` is given, the matching experiment is returned + directly (no prompt), which allows non-interactive/scripted runs. Otherwise, + if a single experiment is found it is returned directly; when multiple experiments are available, the provided ``frontend`` is used to prompt the user to choose one. If no frontend is supplied the default frontend is used. @@ -139,14 +161,16 @@ def _select_experiment(file_path: Path, frontend: Frontend | None = None) -> Exp file_path: Filesystem path to the Python module to inspect. frontend: Optional frontend used to interactively choose an experiment when more than one is discovered. + experiment_name: Optional name of the experiment to select directly, + bypassing the interactive prompt. Returns: ExperimentMetadata: The metadata for the selected experiment. Raises: ValueError: If experiment names are not unique within the module. - SystemExit: If no experiments are found or the user cancels - selection. + SystemExit: If no experiments are found, ``experiment_name`` does not + match any discovered experiment, or the user cancels selection. """ if frontend is None: @@ -161,7 +185,14 @@ def _select_experiment(file_path: Path, frontend: Frontend | None = None) -> Exp msg = f"No @experiment experiments found in {file_path}" raise SystemExit(msg) - if len(experiments) == 1: + if experiment_name is not None: + callable_str_converter = {e.name: e for e in experiments} + if experiment_name not in callable_str_converter: + available = ", ".join(callable_str_converter) + msg = f"No experiment named '{experiment_name}' in {file_path}. Available: {available}" + raise SystemExit(msg) + selected = callable_str_converter[experiment_name] + elif len(experiments) == 1: selected = experiments[0] else: callable_str_converter = {e.name: e for e in experiments} diff --git a/tests/launcher/test_experiments.py b/tests/launcher/test_experiments.py index bfcf182e..874c9fce 100644 --- a/tests/launcher/test_experiments.py +++ b/tests/launcher/test_experiments.py @@ -1,5 +1,8 @@ +from types import ModuleType from unittest.mock import Mock +import pytest + from clabe.launcher import collect_clabe_experiments, experiment, get_experiment_name from clabe.launcher._experiments import _select_experiment from tests import TESTS_ASSETS @@ -33,6 +36,43 @@ def test_select_experiment_multiple_experiments_discovered_and_logs_constant(cap selected.func(launcher) +def test_select_experiment_by_name_skips_prompt() -> None: + mock_frontend = Mock() + + module_path = TESTS_ASSETS / "experiment_import_mocks.py" + selected = _select_experiment(module_path, frontend=mock_frontend, experiment_name="second_experiment") + + assert selected.name == "second_experiment" + mock_frontend.prompt_pick.assert_not_called() + + +def test_select_experiment_by_name_raises_when_not_found() -> None: + module_path = TESTS_ASSETS / "experiment_import_mocks.py" + + with pytest.raises(SystemExit): + _select_experiment(module_path, frontend=Mock(), experiment_name="does_not_exist") + + +def test_collect_clabe_experiments_orders_by_order_field_then_declaration() -> None: + @experiment(name="a") + def a(launcher): ... + + @experiment(name="b") + def b(launcher): ... + + @experiment(name="c", order=-1) + def c(launcher): ... + + # Assignment order (a, then b, then c) mirrors declaration order in a real + # module's namespace; "c" should still sort first via order=-1, while the + # order=0 ties ("a", "b") keep their declaration order. + module = ModuleType("fake_module") + module.a, module.b, module.c = a, b, c + + experiments = list(collect_clabe_experiments(module)) + assert [e.name for e in experiments] == ["c", "a", "b"] + + # --- get_experiment_name ---