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

Signed-off-by: Paul Dittamo <[email protected]>
flyteorg · Oct 16, 2024 · de4e743 · de4e743
1 parent 5b3b33b
commit de4e743
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
@@ -1751,8 +1751,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
@@ -1716,6 +1716,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():