From de4e743d48004127756a58ba74ea616ccc653db9 Mon Sep 17 00:00:00 2001 From: Paul Dittamo <37558497+pvditt@users.noreply.github.com> Date: Wed, 16 Oct 2024 16:07:46 -0700 Subject: [PATCH] fix is_optional_type or not return true for all union types (#2824) Signed-off-by: Paul Dittamo --- flytekit/core/type_engine.py | 3 +-- tests/flytekit/unit/core/test_type_engine.py | 2 ++ 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index c4bf7c08e0..22118fbaf1 100644 --- 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 index dff8a708d2..c5f2dbced0 100644 --- 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():