fix is_optional_type or not return true for all union types (#2824) (#…

…2825) Signed-off-by: Paul Dittamo <[email protected]> Co-authored-by: Paul Dittamo <[email protected]>
flyteorg · Oct 16, 2024 · 1e7306c · 1e7306c
1 parent 54923af
commit 1e7306c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
@@ -1525,8 +1525,7 @@ def __init__(self):
 
     @staticmethod
     def is_optional_type(t: Type) -> bool:
-        """Return True if `t` is a Union or Optional type."""
-        return _is_union_type(t) or type(None) in get_args(t)
+        return _is_union_type(t) and type(None) in get_args(t)
 
     @staticmethod
     def get_sub_type_in_optional(t: Type[T]) -> Type[T]:

diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py
@@ -1707,6 +1707,8 @@ def test_union_transformer():
     assert not UnionTransformer.is_optional_type(str)
     assert UnionTransformer.get_sub_type_in_optional(typing.Optional[int]) == int
     assert UnionTransformer.get_sub_type_in_optional(int | None) == int
+    assert not UnionTransformer.is_optional_type(typing.Union[int, str])
+    assert UnionTransformer.is_optional_type(typing.Union[int, None])
 
 
 def test_union_guess_type():