RF: Strip out template fetching to separate function

Author: mgxd

smriprep/interfaces/templateflow.py

@@ -108,34 +108,9 @@ def _run_interface(self, runtime):
         if isdefined(self.inputs.cohort):
             specs['cohort'] = self.inputs.cohort
 
-        name = self.inputs.template.strip(':').split(':', 1)
-        if len(name) > 1:
-            specs.update(
-                {
-                    k: v
-                    for modifier in name[1].split(':')
-                    for k, v in [tuple(modifier.split('-'))]
-                    if k not in specs
-                }
-            )
-
-        if specs['resolution'] and not isinstance(specs['resolution'], list):
-            specs['resolution'] = [specs['resolution']]
-
-        available_resolutions = tf.TF_LAYOUT.get_resolutions(template=name[0])
-        if specs['resolution'] and not set(specs['resolution']) & set(available_resolutions):
-            fallback_res = available_resolutions[0] if available_resolutions else None
-            LOGGER.warning(
-                f"Template {name[0]} does not have resolution(s): {specs['resolution']}."
-                f"Falling back to resolution: {fallback_res}."
-            )
-            specs['resolution'] = fallback_res
-
-        self._results['t1w_file'] = tf.get(name[0], desc=None, suffix='T1w', **specs)
-
-        self._results['brain_mask'] = tf.get(
-            name[0], desc='brain', suffix='mask', **specs
-        ) or tf.get(name[0], label='brain', suffix='mask', **specs)
+        files = fetch_template_files(self.inputs.template, specs)
+        self._results['t1w_file'] = files['t1w']
+        self._results['brain_mask'] = files['mask']
         return runtime
 
 
@@ -186,3 +161,39 @@ def _run_interface(self, runtime):
                 descsplit = desc.split('-')
                 self._results['spec'][descsplit[0]] = descsplit[1]
         return runtime
+
+
+def fetch_template_files(template: str, specs: dict | None = None) -> dict:
+    if specs is None:
+        specs = {}
+
+    name = template.strip(':').split(':', 1)
+    if len(name) > 1:
+        specs.update(
+            {
+                k: v
+                for modifier in name[1].split(':')
+                for k, v in [tuple(modifier.split('-'))]
+                if k not in specs
+            }
+        )
+
+    if specs['resolution'] and not isinstance(specs['resolution'], list):
+        specs['resolution'] = [specs['resolution']]
+
+    available_resolutions = tf.TF_LAYOUT.get_resolutions(template=name[0])
+    if specs['resolution'] and not set(specs['resolution']) & set(available_resolutions):
+        fallback_res = available_resolutions[0] if available_resolutions else None
+        LOGGER.warning(
+            f"Template {name[0]} does not have resolution(s): {specs['resolution']}."
+            f"Falling back to resolution: {fallback_res}."
+        )
+        specs['resolution'] = fallback_res
+
+    files = {}
+    files['t1w'] = tf.get(name[0], desc=None, suffix='T1w', **specs)
+    files['mask'] = (
+        tf.get(name[0], desc='brain', suffix='mask', **specs)
+        or tf.get(name[0], label='brain', suffix='mask', **specs)
+    )
+    return files

smriprep/utils/misc.py

@@ -21,12 +21,6 @@
 #     https://www.nipreps.org/community/licensing/
 #
 """Self-contained utilities to be used within Function nodes."""
-from pathlib import Path
-import typing as ty
-
-if ty.TYPE_CHECKING:
-    from niworkflows.utils.spaces import Reference
-
 
 def apply_lut(in_dseg, lut, newpath=None):
     """Map the input discrete segmentation to a new label set (lookup table, LUT)."""
@@ -100,25 +94,3 @@ def fs_isRunning(subjects_dir, subject_id, mtime_tol=86400, logger=None):
     if logger:
         logger.warn(f'Removed "IsRunning*" files found under {subj_dir}')
     return subjects_dir
-
-
-def get_template_t1w(template: str, sloppy: bool = False) -> Path:
-    """Query templateflow for the T1w to ensure it is present on the filesystem."""
-    import templateflow.api as tf
-
-    spec = {}
-    _space = template.split(':', 1)
-    if len(_space) > 1:
-        spec['cohort'] = _space[1].replace('cohort-', '')
-    space = _space[0]
-
-    available_res = tf.TF_LAYOUT.get_resolutions(template=space)
-    if sloppy and 2 in available_res:
-        res = 2
-    elif 1 in available_res:
-        res = 1
-    else:
-        res = None
-    spec['resolution'] = res
-
-    return tf.get(space, desc=None, suffix='T1w', **spec)