diff --git a/dagger/dsl/build.py b/dagger/dsl/build.py index 4162a0f..f1c611b 100644 --- a/dagger/dsl/build.py +++ b/dagger/dsl/build.py @@ -154,7 +154,7 @@ def _build_dag_outputs( if isinstance(dag_output, NodeOutputReference): outputs_by_name = {"return_value": dag_output} elif isinstance(dag_output, Mapping): - outputs_by_name = dag_output + outputs_by_name = {f"key_{key}": value for key, value in dag_output.items()} elif dag_output is not None: raise TypeError( f"This DAG returned a value of type {type(dag_output).__name__}. Functions decorated with `dsl.DAG` may only return two types of values: The output of another node or a mapping of [str, the output of another node]" diff --git a/tests/dsl/test_black_box.py b/tests/dsl/test_black_box.py index 70b1e87..ba34c9b 100644 --- a/tests/dsl/test_black_box.py +++ b/tests/dsl/test_black_box.py @@ -376,28 +376,39 @@ def generate_number() -> int: return 1 @dsl.DAG() - def dag(): + def inner_dag(): return { "a": generate_number(), "b": generate_number(), } + @dsl.DAG() + def outer_dag(): + return inner_dag()["a"] + verify_dags_are_equivalent( - dsl.build(dag), + dsl.build(outer_dag), DAG( nodes={ - "generate-number-1": Task( - generate_number.func, - outputs={"return_value": FromReturnValue()}, - ), - "generate-number-2": Task( - generate_number.func, - outputs={"return_value": FromReturnValue()}, + "inner-dag": DAG( + nodes={ + "generate-number-1": Task( + generate_number.func, + outputs={"return_value": FromReturnValue()}, + ), + "generate-number-2": Task( + generate_number.func, + outputs={"return_value": FromReturnValue()}, + ), + }, + outputs={ + "key_a": FromNodeOutput("generate-number-1", "return_value"), + "key_b": FromNodeOutput("generate-number-2", "return_value"), + }, ), }, outputs={ - "a": FromNodeOutput("generate-number-1", "return_value"), - "b": FromNodeOutput("generate-number-2", "return_value"), + "return_value": FromNodeOutput("inner-dag", "key_a"), }, ), ) diff --git a/tests/dsl/test_build.py b/tests/dsl/test_build.py index b0459f3..a69f40d 100644 --- a/tests/dsl/test_build.py +++ b/tests/dsl/test_build.py @@ -257,12 +257,12 @@ def test__build_dag_outputs__when_multiple_outputs_are_returned(): "id-2": "name-2", }, ) == { - "x": FromNodeOutput( + "key_x": FromNodeOutput( "name-1", "output", serializer=AsPickle(), ), - "y": FromNodeOutput( + "key_y": FromNodeOutput( "name-2", "return_value", serializer=AsJSON(indent=1),