ENH: Generalize require for all image types

Author: mgxd

niworkflows/interfaces/bids.py

@@ -261,47 +261,41 @@ class BIDSDataGrabber(SimpleInterface):
     input_spec = _BIDSDataGrabberInputSpec
     output_spec = _BIDSDataGrabberOutputSpec
 
-    def __init__(self, *args, **kwargs):
-        anat_only = kwargs.pop('anat_only', None)
-        anat_derivatives = kwargs.pop('anat_derivatives', None)
-        require_t1w = kwargs.pop('require_t1w', True)
-        require_funcs = kwargs.pop('require_funcs', True)
-        require_pet = kwargs.pop('require_pet', False)
-        super().__init__(*args, **kwargs)
-
+    _image_types = {'asl', 'bold', 'dwi', 'flair', 'fmap', 'pet', 'roi', 'sbref', 't1w', 't2w'}
+
+    def __init__(
+        self,
+        anat_only=False,
+        anat_derivatives=None,
+        *args,
+        **kwargs,
+    ):
+        default_required = {'t1w', 'bold'}
+
+        for imtype in self._image_types:
+            kwarg = f'require_{imtype}'
+            val = kwargs.pop(kwarg, True if imtype in default_required else False)
+            setattr(self, f'_{kwarg}', val)
+
+        self._require_t1w = self._require_t1w and anat_derivatives is None
         if anat_only:
-            self._require_funcs = False
+            self._require_bold = False
             self._require_pet = False
-        else:
-            self._require_funcs = False if require_pet else require_funcs
-            self._require_pet = require_pet
 
-        self._require_t1w = require_t1w and anat_derivatives is None
+        super().__init__(*args, **kwargs)
 
     def _run_interface(self, runtime):
         bids_dict = self.inputs.subject_data
 
         self._results['out_dict'] = bids_dict
         self._results.update(bids_dict)
 
-        if self._require_t1w and not bids_dict['t1w']:
-            raise FileNotFoundError(
-                f'No T1w images found for subject sub-{self.inputs.subject_id}'
-            )
-        
-        if self._require_funcs and not bids_dict['bold']:
-            raise FileNotFoundError(
-                f'No functional images found for subject sub-{self.inputs.subject_id}'
-            )
-        
-        if self._require_pet and not bids_dict.get('pet'):
-            raise FileNotFoundError(
-                f'No PET images found for subject sub-{self.inputs.subject_id}'
-            )
-
-        for imtype in ['t2w', 'flair', 'fmap', 'sbref', 'roi', 'pet', 'asl']:
+        for imtype in self._image_types:
             if not bids_dict[imtype]:
-                LOGGER.info('No "%s" images found for sub-%s', imtype, self.inputs.subject_id)
+                msg = f'No "{imtype}" images found for sub-{self.inputs.subject_id}'
+                if getattr(self, f'_require_{imtype}', False):  # only raise if image is required
+                    raise FileNotFoundError(msg)
+                LOGGER.info(msg)
 
         return runtime
 

niworkflows/interfaces/tests/test_bids.py

@@ -798,20 +798,21 @@ def test_fsdir_min_version(tmp_path, min_version):
 def test_BIDSDataGrabber():
     x = bintfs.BIDSDataGrabber(anat_only=True)
     assert x._require_t1w is True
-    assert x._require_funcs is False
+    assert x._require_bold is False
 
     x = bintfs.BIDSDataGrabber(anat_only=False, require_t1w=False)
     assert x._require_t1w is False
-    assert x._require_funcs is True
+    assert x._require_bold is True
 
     x = bintfs.BIDSDataGrabber(anat_derivatives='derivatives')
     assert x._require_t1w is False
 
 
-def test_require_func_pet_behavior():
-    from niworkflows.interfaces.bids import BIDSDataGrabber
+def test_require_bold_pet_behavior():
     import pytest
 
+    from niworkflows.interfaces.bids import BIDSDataGrabber
+
     subject_data_pet = {
         't1w': ['t1.nii'],
         'bold': [],
@@ -822,6 +823,7 @@ def test_require_func_pet_behavior():
         'sbref': [],
         'roi': [],
         'asl': [],
+        'dwi': [],
     }
 
     subject_data_bold = {
@@ -834,20 +836,23 @@ def test_require_func_pet_behavior():
         'sbref': [],
         'roi': [],
         'asl': [],
+        'dwi': [],
     }
 
     # PET present, functional not required
-    grabber_pet = BIDSDataGrabber(subject_data=subject_data_pet, subject_id='01', require_pet=True)
-    assert grabber_pet._require_funcs is False
+    grabber_pet = BIDSDataGrabber(
+        subject_data=subject_data_pet, subject_id='01', require_pet=True, require_bold=False
+    )
+    assert grabber_pet._require_bold is False
     assert grabber_pet._require_pet is True
     grabber_pet.run()  # Should succeed without bold data
 
     # PET absent, functional required by default
     grabber_func = BIDSDataGrabber(subject_data=subject_data_bold, subject_id='01')
-    assert grabber_func._require_funcs is True
-    grabber_func.run()  # Should succeed with bold data
+    assert grabber_func._require_bold is True
+    assert grabber_func.run().outputs.bold == 'bold.nii'
 
     # Fail when bold is required but missing
     grabber_fail = BIDSDataGrabber(subject_data=subject_data_pet, subject_id='01')
-    with pytest.raises(FileNotFoundError, match='No functional images found'):
-        grabber_fail.run()
+    with pytest.raises(FileNotFoundError, match='No "bold" images found'):
+        grabber_fail.run()