Merge pull request #565 from ghisvail/feat/run-task-dict

Add support for arbitrary mappings in a FunctionTask

pydra/engine/task.py

@@ -198,6 +198,8 @@ def _run_task(self):
             self.output_ = {output_names[0]: output}
         elif isinstance(output, tuple) and len(output_names) == len(output):
             self.output_ = dict(zip(output_names, output))
+        elif isinstance(output, dict):
+            self.output_ = {key: output.get(key, None) for key in output_names}
         else:
             raise RuntimeError(
                 f"expected {len(self.output_spec.fields)} elements, "

pydra/engine/tests/test_task.py

@@ -114,6 +114,27 @@ def testfunc(
     ]
 
 
+def test_annotated_func_dictreturn(use_validator):
+    """Test mapping from returned dictionary to output spec."""
+
+    @mark.task
+    @mark.annotate({"return": {"sum": int, "mul": int}})
+    def testfunc(a: int, b: int):
+        return dict(sum=a + b, diff=a - b)
+
+    task = testfunc(a=2, b=3)
+    result = task()
+
+    # Part of the annotation and returned, should be exposed to output.
+    assert result.output.sum == 5
+
+    # Part of the annotation but not returned, should be coalesced to None.
+    assert result.output.mul is None
+
+    # Not part of the annotation, should be discarded.
+    assert not hasattr(result.output, "diff")
+
+
 def test_annotated_func_multreturn(use_validator):
     """the function has two elements in the return statement"""