Extract BIDS discovery functions from CLI modules

Adds discover_anatomical(), discover_functional(), and
discover_derivative_runs() to the bids package, removing BIDS-specific
iteration logic (entity extraction, row-to-path conversion, DataFrame
filtering/grouping) from CLI modules.

CLI modules now call discovery functions and receive structured
NamedTuples (AnatomicalRun, FunctionalRun, DerivativeRun) instead
of manually iterating DataFrames.

Author: nx10

src/rbc/bids/__init__.py

@@ -24,17 +24,25 @@
     extract_entities,
     parse_bids_name,
 )
-from rbc.bids.anatomical import export_anatomical
+from rbc.bids.anatomical import AnatomicalRun, discover_anatomical, export_anatomical
 from rbc.bids.builder import Bids
-from rbc.bids.functional import FunctionalInputs, export_functional, resolve_functional
+from rbc.bids.functional import (
+    FunctionalInputs,
+    FunctionalRun,
+    discover_functional,
+    export_functional,
+    resolve_functional,
+)
 from rbc.bids.metrics import MetricsInputs, export_metrics, resolve_metrics
 from rbc.bids.qc import QCInputs, export_qc, resolve_qc
 from rbc.bids.query import find_file, find_files, get_extra_entity, load_table
 from rbc.bids.session import (
     ANAT_GROUP_ENTITIES,
     FUNC_GROUP_ENTITIES,
     SUB_SES_QUERY,
+    DerivativeRun,
     SessionTables,
+    discover_derivative_runs,
     iter_session_files,
     load_session,
 )
@@ -45,13 +53,16 @@
     "FUNC_GROUP_ENTITIES",
     "SUB_SES_QUERY",
     "_STANDARD_ENTITIES",
+    "AnatomicalRun",
     "BIDSFile",
     "Bids",
     "BidsEntities",
     "Datatype",
+    "DerivativeRun",
     "EntityKwargs",
     "Extension",
     "FunctionalInputs",
+    "FunctionalRun",
     "MetricsInputs",
     "Modality",
     "QCInputs",
@@ -63,6 +74,9 @@
     "bids_path",
     "bids_path_from_entities",
     "bids_safe_label",
+    "discover_anatomical",
+    "discover_derivative_runs",
+    "discover_functional",
     "export_anatomical",
     "export_functional",
     "export_metrics",

src/rbc/bids/anatomical.py

@@ -1,16 +1,56 @@
-"""BIDS export for the anatomical workflow."""
+"""BIDS discovery and export for the anatomical workflow."""
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from pathlib import Path
+from typing import TYPE_CHECKING, NamedTuple
 
-from rbc.bids import Suffix, TemplateSpace
+import polars as pl
+
+from rbc.bids import Suffix, TemplateSpace, extract_entities
+from rbc.bids.session import ANAT_GROUP_ENTITIES, SessionTables
 
 if TYPE_CHECKING:
-    from rbc.bids import Bids
+    from collections.abc import Iterator
+
+    from rbc.bids import Bids, EntityKwargs
     from rbc.workflows.anatomical import AnatomicalOutputs
 
 
+class AnatomicalRun(NamedTuple):
+    """A single anatomical run discovered from a BIDS session.
+
+    Attributes:
+        path: Path to the T1w NIfTI file.
+        entities: BIDS entities for this run (run, acq, rec, echo).
+    """
+
+    path: Path
+    entities: EntityKwargs
+
+
+def discover_anatomical(session: SessionTables) -> Iterator[AnatomicalRun]:
+    """Discover T1w runs in a session's anatomical data.
+
+    Filters for T1w files and groups by anatomical entities, yielding
+    one :class:`AnatomicalRun` per group.
+
+    Args:
+        session: Session tables from :func:`~rbc.bids.session.load_session`.
+
+    Yields:
+        An :class:`AnatomicalRun` for each T1w group.
+    """
+    for _, anat_df in session.anat.filter(pl.col("suffix") == "T1w").group_by(
+        ANAT_GROUP_ENTITIES, maintain_order=True
+    ):
+        row = anat_df.row(0, named=True)
+        yield AnatomicalRun(
+            path=Path(row["root"]) / row["path"],
+            entities=extract_entities(row, ["run", "acq", "rec", "echo"]),
+        )
+
+
 def export_anatomical(anat: Bids, outputs: AnatomicalOutputs) -> None:
     """Export anatomical workflow outputs to BIDS-named derivatives.
 

src/rbc/bids/functional.py

@@ -1,21 +1,60 @@
-"""BIDS export and resolve for the functional workflow."""
+"""BIDS discovery, resolve, and export for the functional workflow."""
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, TypedDict
+from pathlib import Path
+from typing import TYPE_CHECKING, NamedTuple, TypedDict
 
-from rbc.bids import Suffix, TemplateSpace, bids_safe_label
+import polars as pl
 
-if TYPE_CHECKING:
-    from collections.abc import Sequence
-    from pathlib import Path
+from rbc.bids import Suffix, TemplateSpace, bids_safe_label, extract_entities
+from rbc.bids.session import FUNC_GROUP_ENTITIES, SessionTables, iter_session_files
 
-    import polars as pl
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Sequence
 
-    from rbc.bids import Bids
+    from rbc.bids import Bids, EntityKwargs
     from rbc.workflows.functional import FunctionalOutputs
 
 
+class FunctionalRun(NamedTuple):
+    """A single functional run discovered from a BIDS session.
+
+    Attributes:
+        path: Path to the BOLD NIfTI file.
+        entities: BIDS entities for this run (task, run, acq, rec, dir, echo).
+        anat_df: Matched anatomical DataFrame for this run.
+    """
+
+    path: Path
+    entities: EntityKwargs
+    anat_df: pl.DataFrame
+
+
+def discover_functional(session: SessionTables) -> Iterator[FunctionalRun]:
+    """Discover BOLD runs in a session, paired with matched anatomical data.
+
+    Iterates via :func:`~rbc.bids.session.iter_session_files`, filters for
+    raw (unprocessed) BOLD files, and extracts functional entities.
+
+    Args:
+        session: Session tables from :func:`~rbc.bids.session.load_session`.
+
+    Yields:
+        A :class:`FunctionalRun` for each BOLD group.
+    """
+    for func_df, anat_df in iter_session_files(session, groupby=FUNC_GROUP_ENTITIES):
+        func_df = func_df.filter(pl.col("desc").is_null())
+        row = func_df.filter(suffix="bold").row(0, named=True)
+        yield FunctionalRun(
+            path=Path(row["root"]) / row["path"],
+            entities=extract_entities(
+                row, ["task", "run", "acq", "rec", "dir", "echo"]
+            ),
+            anat_df=anat_df,
+        )
+
+
 class FunctionalInputs(TypedDict):
     """Resolved anatomical inputs for the functional workflow."""
 

src/rbc/bids/session.py

@@ -10,9 +10,13 @@
 
 import polars as pl
 
+from rbc.bids import extract_entities
+
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
 
+    from rbc.bids import EntityKwargs
+
 SUB_SES_QUERY = ("sub", "ses")
 
 ANAT_GROUP_ENTITIES = ("run", "acq", "suffix", "part", "echo", "ce", "rec", "inv")
@@ -127,3 +131,37 @@ def iter_session_files(
             runs_correspond=has_anat_runs and len(anat_runs) == len(func_runs),
         )
         yield func_group, anat_subset
+
+
+_FUNC_ENTITY_KEYS = ("task", "run", "acq", "rec", "dir", "echo")
+
+
+class DerivativeRun(NamedTuple):
+    """A single functional run discovered from derivative data.
+
+    Attributes:
+        entities: BIDS entities for this run (task, run, acq, rec, dir, echo).
+    """
+
+    entities: EntityKwargs
+
+
+def discover_derivative_runs(
+    group: pl.DataFrame,
+) -> Iterator[DerivativeRun]:
+    """Discover functional runs within a sub/ses derivative group.
+
+    Groups by :data:`FUNC_GROUP_ENTITIES` and extracts standard functional
+    entities from each group.
+
+    Args:
+        group: DataFrame of derivative BOLD runs for a single sub/ses.
+
+    Yields:
+        A :class:`DerivativeRun` for each functional run group.
+    """
+    for _, run_group in group.group_by(FUNC_GROUP_ENTITIES):
+        row = run_group.row(0, named=True)
+        yield DerivativeRun(
+            entities=extract_entities(row, list(_FUNC_ENTITY_KEYS)),
+        )

src/rbc/cli/all.py

@@ -8,25 +8,17 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING, Literal
 
 import polars as pl
 from tqdm import tqdm
 
-from rbc.bids import (
-    ANAT_GROUP_ENTITIES,
-    FUNC_GROUP_ENTITIES,
-    SUB_SES_QUERY,
-    Datatype,
-    extract_entities,
-    load_table,
-)
-from rbc.bids.anatomical import export_anatomical
-from rbc.bids.functional import export_functional
+from rbc.bids import SUB_SES_QUERY, Datatype, load_table
+from rbc.bids.anatomical import discover_anatomical, export_anatomical
+from rbc.bids.functional import discover_functional, export_functional
 from rbc.bids.metrics import export_metrics
 from rbc.bids.qc import export_qc
-from rbc.bids.session import iter_session_files, load_session
+from rbc.bids.session import load_session
 from rbc.cli import _DEFAULT_ENV_VARS
 from rbc.cli.base import BaseArgs, _validate_atlas, _validate_positive, _validate_task
 from rbc.context import RunContext
@@ -109,32 +101,22 @@ def main(args: AllArgs) -> int:
         session = load_session(sub_ses_group, pipe_ctx.sub, pipe_ctx.ses)
 
         # --- Anatomical (once per session, first T1w) ---
-        for _, anat_df in session.anat.filter(pl.col("suffix") == "T1w").group_by(
-            ANAT_GROUP_ENTITIES, maintain_order=True
-        ):
-            anat_row = anat_df.filter(suffix="T1w").row(0, named=True)
-            t1w_fpath = Path(anat_row["root"]) / anat_row["path"]
-            ents = extract_entities(anat_row, ["run", "acq", "rec", "echo"])
-            ctx.logger.info(f"Anatomical: {t1w_fpath}")
+        for anat_run in discover_anatomical(session):
+            ctx.logger.info(f"Anatomical: {anat_run.path}")
 
-            anat_outputs = anatomical_preprocess(in_t1w=t1w_fpath)
+            anat_outputs = anatomical_preprocess(in_t1w=anat_run.path)
 
-            anat = pipe_ctx.bids(datatype=Datatype.ANAT, entities=ents)
+            anat = pipe_ctx.bids(datatype=Datatype.ANAT, entities=anat_run.entities)
             export_anatomical(anat, anat_outputs)
 
         # --- Functional + Metrics + QC (per BOLD run) ---
-        for func_df, _anat_df in iter_session_files(
-            session, groupby=FUNC_GROUP_ENTITIES
-        ):
-            row = func_df.row(0, named=True)
-            bold_fpath = Path(row["root"]) / row["path"]
-            ents = extract_entities(row, ["task", "run", "acq", "rec", "dir", "echo"])
-            ctx.logger.info(f"Functional: {bold_fpath}")
+        for func_run in discover_functional(session):
+            ctx.logger.info(f"Functional: {func_run.path}")
 
-            func_metadata = FunctionalMetadata.load(bold_fpath, tr_override=args.tr)
+            func_metadata = FunctionalMetadata.load(func_run.path, tr_override=args.tr)
 
             func_outputs = functional_preprocess(
-                in_bold=bold_fpath,
+                in_bold=func_run.path,
                 t1w_brain=anat_outputs.brain,
                 wm_bbr_mask=anat_outputs.wm_bbr_mask,
                 brain_mask=anat_outputs.brain_mask,
@@ -146,14 +128,16 @@ def main(args: AllArgs) -> int:
                 regressor_set=args.regressor,
             )
 
-            func = pipe_ctx.bids(datatype=Datatype.FUNC, entities=ents)
+            func = pipe_ctx.bids(datatype=Datatype.FUNC, entities=func_run.entities)
             mni = export_functional(func, func_outputs, regressors=args.regressor)
 
             # --- Metrics ---
             for regressor in args.regressor:
+                task = func_run.entities.get("task", "")
+                run = func_run.entities.get("run", 0)
                 ctx.logger.info(
-                    f"Metrics: sub-{pipe_ctx.sub} task-{ents.get('task', '')} "
-                    f"run-{ents.get('run', 0)} regressor-{regressor}"
+                    f"Metrics: sub-{pipe_ctx.sub} task-{task} "
+                    f"run-{run} regressor-{regressor}"
                 )
                 metrics_outputs = metrics_pipeline(
                     regressed_bold=func_outputs.regressed_bold[regressor],
@@ -172,10 +156,7 @@ def main(args: AllArgs) -> int:
                 )
 
             # --- QC ---
-            ctx.logger.info(
-                f"QC: sub-{pipe_ctx.sub} "
-                f"task-{ents.get('task', '')} run-{ents.get('run', 0)}"
-            )
+            ctx.logger.info(f"QC: sub-{pipe_ctx.sub} task-{task} run-{run}")
             qc_outputs = qc_pipeline(
                 template_bold=func_outputs.template_bold,
                 cleaned_bold=func_outputs.cleaned_bold,
@@ -187,19 +168,16 @@ def main(args: AllArgs) -> int:
                 template_brain_mask=func_outputs.template_brain_mask,
                 sub=pipe_ctx.sub,
                 ses=pipe_ctx.ses or "",
-                task=ents.get("task", ""),
-                run=ents.get("run", 0),
+                task=func_run.entities.get("task", ""),
+                run=func_run.entities.get("run", 0),
                 start_tr=args.start_tr,
                 regressor_set=args.regressor,
             )
 
             export_qc(mni, qc_outputs, regressors=args.regressor)
 
             status = "PASSED" if qc_outputs.passed else "FAILED"
-            ctx.logger.info(
-                f"QC {status} for sub-{pipe_ctx.sub} task-{ents.get('task', '')} "
-                f"run-{ents.get('run', 0)}"
-            )
+            ctx.logger.info(f"QC {status} for sub-{pipe_ctx.sub} task-{task} run-{run}")
         pipe_ctx.ensure_dataset_description()
 
     ctx.logger.info("RBC full pipeline complete")