diff --git a/pydra/engine/task.py b/pydra/engine/task.py index 06f5133530..c6125fbadd 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"""