From b89f9a2453c3ad36ad28b9efba7320676a7ba44e Mon Sep 17 00:00:00 2001
From: Ghislain Vaillant <ghislain.vaillant@icm-institute.org>
Date: Tue, 16 Aug 2022 15:21:57 +0200
Subject: [PATCH] ENH: Add support for mappings in a FunctionTask

---
 pydra/engine/task.py            |  2 ++
 pydra/engine/tests/test_task.py | 21 +++++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/pydra/engine/task.py b/pydra/engine/task.py
index 135af50796..a05e26a457 100644
--- a/pydra/engine/task.py
+++ b/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, "
diff --git a/pydra/engine/tests/test_task.py b/pydra/engine/tests/test_task.py
index 7ef386f348..2762bcf950 100644
--- a/pydra/engine/tests/test_task.py
+++ b/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"""