accepting None as an output from function that is expected to return multiple values

Author: djarecka

pydra/engine/task.py

@@ -200,17 +200,20 @@ def _run_task(self):
         self.output_ = None
         output = cp.loads(self.inputs._func)(**inputs)
         output_names = [el[0] for el in self.output_spec.fields]
-        self.output_ = {}
-        if len(output_names) > 1:
-            if len(output_names) == len(output):
-                self.output_ = dict(zip(output_names, output))
+        if output is None:
+            self.output_ = dict((nm, None) for nm in output_names)
+        else:
+            if len(output_names) == 1:
+                # if only one element in the fields, everything should be returned together
+                self.output_ = {output_names[0]: output}
             else:
-                raise Exception(
-                    f"expected {len(self.output_spec.fields)} elements, "
-                    f"but {len(output)} were returned"
-                )
-        else:  # if only one element in the fields, everything should be returned together
-            self.output_[output_names[0]] = output
+                if isinstance(output, tuple) and len(output_names) == len(output):
+                    self.output_ = dict(zip(output_names, output))
+                else:
+                    raise Exception(
+                        f"expected {len(self.output_spec.fields)} elements, "
+                        f"but {output} were returned"
+                    )
 
 
 class ShellCommandTask(TaskBase):

pydra/engine/tests/test_task.py

@@ -309,7 +309,7 @@ def raise_exception(c, d):
     assert pytest.raises(Exception, bad_funk)
 
 
-def test_result_none():
+def test_result_none_1():
     """ checking if None is properly returned as the result"""
 
     @mark.task
@@ -321,6 +321,19 @@ def fun_none(x):
     assert res.output.out is None
 
 
+def test_result_none_2():
+    """ checking if None is properly set for all outputs """
+
+    @mark.task
+    def fun_none(x) -> (ty.Any, ty.Any):
+        return None
+
+    task = fun_none(name="none", x=3)
+    res = task()
+    assert res.output.out1 is None
+    assert res.output.out2 is None
+
+
 def test_audit_prov(tmpdir):
     @mark.task
     def testfunc(a: int, b: float = 0.1) -> ty.NamedTuple("Output", [("out", float)]):