Merge pull request #955 from oesteban/enh/niuSplitSqueeze

Squeeze option in utility.Split

Author: chrisgorgo

nipype/interfaces/tests/test_auto_Split.py

@@ -10,6 +10,8 @@ def test_Split_inputs():
     ),
     splits=dict(mandatory=True,
     ),
+    squeeze=dict(usedefault=True,
+    ),
     )
     inputs = Split.input_spec()
 

nipype/interfaces/tests/test_utility.py

@@ -113,3 +113,28 @@ def test_function_with_imports():
     finally:
         os.chdir(origdir)
         shutil.rmtree(tempdir)
+
+
+def test_split():
+    tempdir = os.path.realpath(mkdtemp())
+    origdir = os.getcwd()
+    os.chdir(tempdir)
+
+    try:
+        node = pe.Node(utility.Split(inlist=range(4),
+                                     splits=[1, 3]),
+                       name='split_squeeze')
+        res = node.run()
+        yield assert_equal, res.outputs.out1, [0]
+        yield assert_equal, res.outputs.out2, [1, 2, 3]
+
+        node = pe.Node(utility.Split(inlist=range(4),
+                                     splits=[1, 3],
+                                     squeeze=True),
+                       name='split_squeeze')
+        res = node.run()
+        yield assert_equal, res.outputs.out1, 0
+        yield assert_equal, res.outputs.out2, [1, 2, 3]
+    finally:
+        os.chdir(origdir)
+        shutil.rmtree(tempdir)

nipype/interfaces/utility.py

@@ -252,6 +252,8 @@ class SplitInputSpec(BaseInterfaceInputSpec):
                   desc='list of values to split')
     splits = traits.List(traits.Int, mandatory=True,
                   desc='Number of outputs in each split - should add to number of inputs')
+    squeeze = traits.Bool(False, usedefault=True,
+                          desc='unfold one-element splits removing the list')
 
 
 class Split(IOBase):
@@ -290,7 +292,10 @@ def _list_outputs(self):
             splits.extend(self.inputs.splits)
             splits = np.cumsum(splits)
             for i in range(len(splits) - 1):
-                outputs['out%d' % (i + 1)] = np.array(self.inputs.inlist)[splits[i]:splits[i + 1]].tolist()
+                val = np.array(self.inputs.inlist)[splits[i]:splits[i + 1]].tolist()
+                if self.inputs.squeeze and len(val) == 1:
+                    val = val[0]
+                outputs['out%d' % (i + 1)] = val
         return outputs