fix: apply review comments from @effigies, add parameterized tests

Author: oesteban

niworkflows/interfaces/nibabel.py

@@ -101,9 +101,6 @@ def _run_interface(self, runtime):
 
 class _SplitSeriesInputSpec(BaseInterfaceInputSpec):
     in_file = File(exists=True, mandatory=True, desc="input 4d image")
-    allow_3D = traits.Bool(
-        False, usedefault=True, desc="do not fail if a 3D volume is passed in"
-    )
 
 
 class _SplitSeriesOutputSpec(TraitedSpec):
@@ -117,34 +114,20 @@ class SplitSeries(SimpleInterface):
     output_spec = _SplitSeriesOutputSpec
 
     def _run_interface(self, runtime):
-        filenii = nb.squeeze_image(nb.load(self.inputs.in_file))
-        filenii = filenii.__class__(
-            np.squeeze(filenii.dataobj), filenii.affine, filenii.header
-        )
-        ndim = filenii.dataobj.ndim
-        if ndim != 4:
-            if self.inputs.allow_3D and ndim == 3:
-                out_file = str(
-                    Path(
-                        fname_presuffix(self.inputs.in_file, suffix=f"_idx-000")
-                    ).absolute()
-                )
-                self._results["out_files"] = out_file
-                filenii.to_filename(out_file)
-                return runtime
-            raise RuntimeError(
-                f"Input image <{self.inputs.in_file}> is {ndim}D "
-                f"({'x'.join(['%d' % s for s in filenii.shape])})."
-            )
+        in_file = self.inputs.in_file
+        img = nb.load(in_file)
+        extra_dims = tuple(dim for dim in img.shape[3:] if dim > 1) or (1,)
+        if len(extra_dims) != 1:
+            raise ValueError(f"Invalid shape {'x'.join(str(s) for s in img.shape)}")
+        img = img.__class__(img.dataobj.reshape(img.shape[:3] + extra_dims),
+                            img.affine, img.header)
 
-        files_3d = nb.four_to_three(filenii)
         self._results["out_files"] = []
-        in_file = self.inputs.in_file
-        for i, file_3d in enumerate(files_3d):
+        for i, img_3d in enumerate(nb.four_to_three(img)):
             out_file = str(
                 Path(fname_presuffix(in_file, suffix=f"_idx-{i:03}")).absolute()
             )
-            file_3d.to_filename(out_file)
+            img_3d.to_filename(out_file)
             self._results["out_files"].append(out_file)
 
         return runtime

niworkflows/interfaces/tests/test_nibabel.py

@@ -75,61 +75,35 @@ def test_ApplyMask(tmp_path):
         ApplyMask(in_file=str(in_file4d), in_mask=str(in_mask), threshold=0.4).run()
 
 
-def test_SplitSeries(tmp_path):
+@pytest.mark.parametrize("shape,exp_n", [
+    ((20, 20, 20, 15), 15),
+    ((20, 20, 20), 1),
+    ((20, 20, 20, 1), 1),
+    ((20, 20, 20, 1, 3), 3),
+    ((20, 20, 20, 3, 1), 3),
+    ((20, 20, 20, 1, 3, 3), -1),
+    ((20, 1, 20, 15), 15),
+    ((20, 1, 20), 1),
+    ((20, 1, 20, 1), 1),
+    ((20, 1, 20, 1, 3), 3),
+    ((20, 1, 20, 3, 1), 3),
+    ((20, 1, 20, 1, 3, 3), -1),
+])
+def test_SplitSeries(tmp_path, shape, exp_n):
     """Test 4-to-3 NIfTI split interface."""
     os.chdir(tmp_path)
 
-    # Test the 4D
-    data = np.ones((20, 20, 20, 15), dtype=float)
-    in_file = tmp_path / "input4D.nii.gz"
-    nb.Nifti1Image(data, np.eye(4), None).to_filename(str(in_file))
-
-    split = SplitSeries(in_file=str(in_file)).run()
-    assert len(split.outputs.out_files) == 15
-
-    # Test the 3D
-    in_file = tmp_path / "input3D.nii.gz"
-    nb.Nifti1Image(np.ones((20, 20, 20), dtype=float), np.eye(4), None).to_filename(
-        str(in_file)
-    )
-
-    with pytest.raises(RuntimeError):
-        SplitSeries(in_file=str(in_file)).run()
-
-    split = SplitSeries(in_file=str(in_file), allow_3D=True).run()
-    assert isinstance(split.outputs.out_files, str)
-
-    # Test the 3D
-    in_file = tmp_path / "input3D.nii.gz"
-    nb.Nifti1Image(np.ones((20, 20, 20, 1), dtype=float), np.eye(4), None).to_filename(
-        str(in_file)
-    )
-
-    with pytest.raises(RuntimeError):
-        SplitSeries(in_file=str(in_file)).run()
-
-    split = SplitSeries(in_file=str(in_file), allow_3D=True).run()
-    assert isinstance(split.outputs.out_files, str)
-
-    # Test the 5D
-    in_file = tmp_path / "input5D.nii.gz"
-    nb.Nifti1Image(
-        np.ones((20, 20, 20, 2, 2), dtype=float), np.eye(4), None
-    ).to_filename(str(in_file))
-
-    with pytest.raises(RuntimeError):
-        SplitSeries(in_file=str(in_file)).run()
-
-    with pytest.raises(RuntimeError):
-        SplitSeries(in_file=str(in_file), allow_3D=True).run()
-
-    # Test splitting ANTs warpfields
-    data = np.ones((20, 20, 20, 1, 3), dtype=float)
-    in_file = tmp_path / "warpfield.nii.gz"
-    nb.Nifti1Image(data, np.eye(4), None).to_filename(str(in_file))
-
-    split = SplitSeries(in_file=str(in_file)).run()
-    assert len(split.outputs.out_files) == 3
+    in_file = str(tmp_path / "input.nii.gz")
+    nb.Nifti1Image(np.ones(shape, dtype=float), np.eye(4), None).to_filename(in_file)
+
+    _interface = SplitSeries(in_file=in_file)
+    if exp_n > 0:
+        split = _interface.run()
+        n = int(isinstance(split.outputs.out_files, str)) or len(split.outputs.out_files)
+        assert n == exp_n
+    else:
+        with pytest.raises(ValueError):
+            _interface.run()
 
 
 def test_MergeSeries(tmp_path):