Fixes guess type bug in UnionTransformer

Signed-off-by: Ketan Umare <[email protected]>
flyteorg · Jan 31, 2023 · 1a28915 · 1a28915
1 parent 1513773
commit 1a28915
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
@@ -1138,7 +1138,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
 
     def guess_python_type(self, literal_type: LiteralType) -> type:
         if literal_type.union_type is not None:
-            return typing.Union[tuple(TypeEngine.guess_python_type(v.type) for v in literal_type.union_type.variants)]  # type: ignore
+            return typing.Union[tuple(TypeEngine.guess_python_type(v) for v in literal_type.union_type.variants)]
 
         raise ValueError(f"Union transformer cannot reverse {literal_type}")
 

diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py
@@ -45,7 +45,7 @@
 from flytekit.models.annotation import TypeAnnotation
 from flytekit.models.core.types import BlobType
 from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Void
-from flytekit.models.types import LiteralType, SimpleType, TypeStructure
+from flytekit.models.types import LiteralType, SimpleType, TypeStructure, UnionType
 from flytekit.types.directory import TensorboardLogs
 from flytekit.types.directory.types import FlyteDirectory
 from flytekit.types.file import FileExt, JPEGImageFile
@@ -941,6 +941,18 @@ def test_union_transformer():
     assert UnionTransformer.get_sub_type_in_optional(typing.Optional[int]) == int
 
 
+def test_union_guess_type():
+    ut = UnionTransformer()
+    t = ut.guess_python_type(
+        LiteralType(
+            union_type=UnionType(
+                variants=[LiteralType(simple=SimpleType.STRING), LiteralType(simple=SimpleType.INTEGER)]
+            )
+        )
+    )
+    assert t == typing.Union[str, int]
+
+
 def test_union_type_with_annotated():
     pt = typing.Union[
         Annotated[str, FlyteAnnotation({"hello": "world"})], Annotated[int, FlyteAnnotation({"test": 123})]