♻️ SSOT warp_mask_to_template decorator

CPAC/pipeline/cpac_pipeline.py

@@ -17,6 +17,7 @@
 import os
 import sys
 import time
+from typing import Literal
 import csv
 import shutil
 import pickle
@@ -1033,7 +1034,8 @@ def build_anat_preproc_stack(rpool, cfg, pipeline_blocks=None):
     return pipeline_blocks
 
 
-def build_T1w_registration_stack(rpool, cfg, pipeline_blocks=None, space="T1w"):
+def build_T1w_registration_stack(rpool, cfg, pipeline_blocks=None,
+                                 space: Literal["longitudinal", "T1w"] = "T1w"):
 
     if not pipeline_blocks:
         pipeline_blocks = []
@@ -1044,7 +1046,7 @@ def build_T1w_registration_stack(rpool, cfg, pipeline_blocks=None, space="T1w"):
             [register_ANTs_anat_to_template, register_FSL_anat_to_template],
             overwrite_transform_anat_to_template,
             warp_wholeheadT1_to_template,
-            warp_mask_to_template[space]
+            warp_mask_to_template(space)
         ]
 
     if not rpool.check_rpool('desc-restore-brain_T1w'):

CPAC/registration/registration.py

@@ -16,9 +16,9 @@
 # License along with C-PAC. If not, see <https://www.gnu.org/licenses/>.
 # pylint: disable=too-many-lines,ungrouped-imports,wrong-import-order
 from copy import deepcopy
-from typing import Optional
+from typing import Literal, Optional
 from CPAC.pipeline import nipype_pipeline_engine as pe
-from CPAC.pipeline.nodeblock import nodeblock
+from CPAC.pipeline.nodeblock import nodeblock, NodeBlockFunction
 from nipype.interfaces import afni, ants, c3, fsl, utility as util
 from nipype.interfaces.afni import utils as afni_utils
 
@@ -3572,8 +3572,24 @@ def warp_wholeheadT1_to_template(wf, cfg, strat_pool, pipe_num, opt=None):
 
     return (wf, outputs)
 
-def _warp_mask_to_template(space):
+def warp_mask_to_template(space: Literal["longitudinal", "T1w"]) -> NodeBlockFunction:
+    """Get a NodeBlockFunction to transform a mask from ``space`` to template."""
+    @nodeblock(
+        name=f"transform_{space}-mask_to_T1-template",
+        switch=[
+            ["registration_workflows", "anatomical_registration", "run"],
+            ["anatomical_preproc", "run"],
+            ["anatomical_preproc", "brain_extraction", "run"],
+        ],
+        inputs=[
+            (f"space-{space}_desc-brain_mask",
+             f"from-{space}_to-template_mode-image_xfm"),
+            "T1w-template",
+        ],
+        outputs={"space-template_desc-brain_mask": {"Template": "T1w-template"}},
+    )
     def warp_mask_to_template_fxn(wf, cfg, strat_pool, pipe_num, opt=None):
+        """Transform a mask to template space."""
 
         xfm_prov = strat_pool.get_cpac_provenance(
             f'from-{space}_to-template_mode-image_xfm')
@@ -3616,32 +3632,6 @@ def warp_mask_to_template_fxn(wf, cfg, strat_pool, pipe_num, opt=None):
         return wf, outputs
     return warp_mask_to_template_fxn
 
-warp_mask_to_template = {"T1w": nodeblock(
-    name="transform_T1mask_to_T1template",
-    switch=[
-        ["registration_workflows", "anatomical_registration", "run"],
-        ["anatomical_preproc", "run"],
-        ["anatomical_preproc", "brain_extraction", "run"],
-    ],
-    inputs=[
-        ("space-T1w_desc-brain_mask", "from-T1w_to-template_mode-image_xfm"),
-        "T1w-template",
-    ],
-    outputs={"space-template_desc-brain_mask": {"Template": "T1w-template"}},
-)(_warp_mask_to_template("T1w")), "longitudinal": nodeblock(
-    name="transform_longitudinal_mask_to_T1template",
-    switch=[
-        ["registration_workflows", "anatomical_registration", "run"],
-        ["anatomical_preproc", "run"],
-        ["anatomical_preproc", "brain_extraction", "run"],
-    ],
-    inputs=[
-        ("space-longitudinal_desc-brain_mask", "from-longitudinal_to-template_mode-image_xfm"),
-        "T1w-template",
-    ],
-    outputs={"space-template_desc-brain_mask": {"Template": "T1w-template"}},
-)(_warp_mask_to_template("longitudinal"))}
-
 
 @nodeblock(
     name="transform_timeseries_to_T1template",