ENH: Enable Merge(0) to merge a list of lists

effigies · effigies · commit 73f93a2a9192 · 2017-03-21T22:17:25.000-04:00
diff --git a/nipype/interfaces/utility/base.py b/nipype/interfaces/utility/base.py
@@ -128,24 +128,38 @@ class Merge(IOBase):
     def __init__(self, numinputs=0, **inputs):
         super(Merge, self).__init__(**inputs)
         self._numinputs = numinputs
-        add_traits(self.inputs, ['in%d' % (i + 1) for i in range(numinputs)])
+        if numinputs > 0:
+            input_names = ['in%d' % (i + 1) for i in range(numinputs)]
+        elif numinputs == 0:
+            input_names = ['in_lists']
+        else:
+            input_names = []
+        add_traits(self.inputs, input_names)
 
     def _list_outputs(self):
         outputs = self._outputs().get()
         out = []
+
+        if self._numinputs == 0:
+            values = getattr(self.inputs, 'in_lists')
+            if not isdefined(values):
+                return outputs
+        else:
+            getval = lambda idx: getattr(self.inputs, 'in%d' % (idx + 1))
+            values = [getval(idx) for idx in range(self._numinputs)
+                      if isdefined(getval(idx))]
+
         if self.inputs.axis == 'vstack':
-            for idx in range(self._numinputs):
-                value = getattr(self.inputs, 'in%d' % (idx + 1))
-                if isdefined(value):
-                    if isinstance(value, list) and not self.inputs.no_flatten:
-                        out.extend(value)
-                    else:
-                        out.append(value)
+            for value in values:
+                if isinstance(value, list) and not self.inputs.no_flatten:
+                    out.extend(value)
+                else:
+                    out.append(value)
         else:
-            for i in range(len(filename_to_list(self.inputs.in1))):
+            for i in range(len(filename_to_list(values[0]))):
                 out.insert(i, [])
-                for j in range(self._numinputs):
-                    out[i].append(filename_to_list(getattr(self.inputs, 'in%d' % (j + 1)))[i])
+                for value in values:
+                    out[i].append(filename_to_list(value)[i])
         if out:
             outputs['out'] = out
         return outputs
diff --git a/nipype/interfaces/utility/tests/test_base.py b/nipype/interfaces/utility/tests/test_base.py
@@ -55,23 +55,32 @@ def test_split(tmpdir, args, expected):
 @pytest.mark.parametrize("args, kwargs, in_lists, expected", [
         ([3], {}, [0, [1, 2], [3, 4, 5]], [0, 1, 2, 3, 4, 5]),
         ([], {}, None, None),
+        ([], {}, [0, [1, 2], [3, 4, 5]], [0, 1, 2, 3, 4, 5]),
         ([3], {'axis': 'hstack'}, [[0], [1, 2], [3, 4, 5]], [[0, 1, 3]]),
         ([3], {'axis': 'hstack'}, [[0, 1], [2, 3], [4, 5]],
          [[0, 2, 4], [1, 3, 5]]),
         ([3], {'axis': 'hstack'}, [[0, 1], [2, 3], [4, 5]],
          [[0, 2, 4], [1, 3, 5]]),
+        # Note: Merge(0, axis='hstack') would error on run, prior to
+        # in_lists implementation
+        ([0], {'axis': 'hstack'}, [[0], [1, 2], [3, 4, 5]], [[0, 1, 3]]),
+        ([0], {'axis': 'hstack'}, [[0, 1], [2, 3], [4, 5]],
+         [[0, 2, 4], [1, 3, 5]]),
         ])
 def test_merge(tmpdir, args, kwargs, in_lists, expected):
     os.chdir(str(tmpdir))
 
     node = pe.Node(utility.Merge(*args, **kwargs), name='merge')
 
     numinputs = args[0] if args else 0
-    for i in range(1, numinputs + 1):
-        setattr(node.inputs, 'in{:d}'.format(i), in_lists[i - 1])
+    if numinputs == 0 and in_lists:
+        node.inputs.in_lists = in_lists
+    else:
+        for i in range(1, numinputs + 1):
+            setattr(node.inputs, 'in{:d}'.format(i), in_lists[i - 1])
 
     res = node.run()
-    if numinputs == 0:
+    if numinputs == 0 and in_lists is None:
         assert not isdefined(res.outputs.out)
     else:
         assert res.outputs.out == expected
