Use bids entities for filters

Author: kaitj

src/rbc/bids/longitudinal/template.py

@@ -7,7 +7,7 @@
 
 import polars as pl
 
-from rbc.bids import Suffix, bids_safe_label
+from rbc.bids import Datatype, Suffix, bids_safe_label
 
 if TYPE_CHECKING:
     from rbc.bids import Bids
@@ -47,9 +47,9 @@ def discover_template_inputs(
     """
     filtered = df.filter(
         pl.col("ses") != "longitudinal",
-        pl.col("datatype") == "anat",
+        pl.col("datatype") == Datatype.ANAT,
         pl.col("desc") == "brain",
-        pl.col("suffix") == "T1w",
+        pl.col("suffix") == Suffix.T1W,
         # Native-space brains only; the MNI-registered desc-brain T1w that
         # cross-sectional anat also writes would otherwise be picked up as
         # a second input per session, producing duplicate LTA filenames in

src/rbc/bids/session.py

@@ -10,7 +10,7 @@
 
 import polars as pl
 
-from rbc.bids import extract_entities
+from rbc.bids import Datatype, extract_entities
 
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
@@ -50,8 +50,8 @@ def load_session(df: pl.DataFrame, subject: str, session: str | None) -> Session
     base: list[pl.Expr] = [pl.col("sub") == subject]
     if session is not None:
         base.append(pl.col("ses") == session)
-    anat_df = df.filter(pl.all_horizontal([*base, pl.col("datatype") == "anat"]))
-    func_df = df.filter(pl.all_horizontal([*base, pl.col("datatype") == "func"]))
+    anat_df = df.filter(pl.all_horizontal([*base, pl.col("datatype") == Datatype.ANAT]))
+    func_df = df.filter(pl.all_horizontal([*base, pl.col("datatype") == Datatype.FUNC]))
 
     return SessionTables(anat=anat_df, func=func_df if not func_df.is_empty() else None)
 

src/rbc/orchestration/longitudinal/qc.py

@@ -144,12 +144,12 @@ def run(
         anat_q = pipe_ctx.bids(datatype=Datatype.ANAT)
         anat_long_q = anat_q.derive(space="longitudinal")
         anat_long_df = full_df.filter(
-            pl.col("datatype") == "anat",
+            pl.col("datatype") == Datatype.ANAT,
             pl.col("space") == "longitudinal",
         )
 
         func_long_df = full_df.filter(
-            pl.col("datatype") == "func",
+            pl.col("datatype") == Datatype.FUNC,
             pl.col("space") == "longitudinal",
         )
 

src/rbc/orchestration/metrics.py

@@ -9,7 +9,7 @@
 import polars as pl
 from tqdm import tqdm
 
-from rbc.bids import SUB_SES_QUERY, Datatype, TemplateSpace, load_table
+from rbc.bids import SUB_SES_QUERY, Datatype, Suffix, TemplateSpace, load_table
 from rbc.bids.metrics import export_metrics, resolve_metrics
 from rbc.bids.session import discover_derivative_runs
 from rbc.context import RunContext
@@ -109,8 +109,8 @@ def run(
 
     df = filters.apply(
         full_df,
-        pl.col("datatype") == "func",
-        pl.col("suffix") == "bold",
+        pl.col("datatype") == Datatype.FUNC,
+        pl.col("suffix") == Suffix.BOLD,
         pl.col("desc") == "preproc",
         pl.col("space") == TemplateSpace.MNI152NLIN6ASYM,
     )

src/rbc/orchestration/qc.py

@@ -8,7 +8,7 @@
 import polars as pl
 from tqdm import tqdm
 
-from rbc.bids import SUB_SES_QUERY, Datatype, TemplateSpace, load_table
+from rbc.bids import SUB_SES_QUERY, Datatype, Suffix, TemplateSpace, load_table
 from rbc.bids.qc import export_qc, resolve_qc
 from rbc.bids.session import discover_derivative_runs
 from rbc.context import RunContext
@@ -53,8 +53,8 @@ def run(
 
     df = filters.apply(
         full_df,
-        pl.col("datatype") == "func",
-        pl.col("suffix") == "bold",
+        pl.col("datatype") == Datatype.FUNC,
+        pl.col("suffix") == Suffix.BOLD,
         pl.col("desc") == "preproc",
         pl.col("space") == TemplateSpace.MNI152NLIN6ASYM,
     )