Merge pull request #261 from djarecka/fix/result_none

[fix] allowing annotated function to return none (closes #256)

Author: djarecka

pydra/engine/task.py

@@ -199,19 +199,21 @@ def _run_task(self):
         del inputs["_func"]
         self.output_ = None
         output = cp.loads(self.inputs._func)(**inputs)
-        if output is not None:
-            output_names = [el[0] for el in self.output_spec.fields]
-            self.output_ = {}
-            if len(output_names) > 1:
-                if len(output_names) == len(output):
+        output_names = [el[0] for el in self.output_spec.fields]
+        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:
+                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 {len(output)} were returned"
+                        f"but {output} were returned"
                     )
-            else:  # if only one element in the fields, everything should be returned together
-                self.output_[output_names[0]] = output
 
 
 class ShellCommandTask(TaskBase):

pydra/engine/tests/test_task.py

@@ -309,6 +309,31 @@ def raise_exception(c, d):
     assert pytest.raises(Exception, bad_funk)
 
 
+def test_result_none_1():
+    """ checking if None is properly returned as the result"""
+
+    @mark.task
+    def fun_none(x):
+        return None
+
+    task = fun_none(name="none", x=3)
+    res = task()
+    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)]):