Add support dict in @dynamic (#1264)

* add items() Signed-off-by: Kevin Su <[email protected]> * More tests Signed-off-by: Kevin Su <[email protected]> * Fix tests Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Kevin Su <[email protected]>
flyteorg · Oct 31, 2022 · b787849 · b787849
1 parent cef02aa
commit b787849
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py
@@ -464,7 +464,9 @@ def to_literal_model(self):
                 )
             )
         elif self.map:
-            return Literal(map=LiteralMap(literals={k: binding.to_literal_model() for k, binding in self.map.bindings}))
+            return Literal(
+                map=LiteralMap(literals={k: binding.to_literal_model() for k, binding in self.map.bindings.items()})
+            )
 
 
 class Binding(_common.FlyteIdlEntity):

diff --git a/tests/flytekit/unit/core/test_dynamic.py b/tests/flytekit/unit/core/test_dynamic.py
@@ -153,3 +153,25 @@ def wf() -> str:
         return dynamic_wf()
 
     assert wf() == "hello"
+
+
+def test_dynamic_return_dict():
+    @dynamic
+    def t1(v: str) -> typing.Dict[str, str]:
+        return {"a": v}
+
+    @dynamic
+    def t2(v: str) -> typing.Dict[str, typing.Dict[str, str]]:
+        return {"a": {"b": v}}
+
+    @dynamic
+    def t3(v: str) -> (str, typing.Dict[str, typing.Dict[str, str]]):
+        return v, {"a": {"b": v}}
+
+    @workflow
+    def wf():
+        t1(v="a")
+        t2(v="b")
+        t3(v="c")
+
+    wf()