From 2574c9f8735b78fa8eed801f37a048879c99cae0 Mon Sep 17 00:00:00 2001 From: Daniel Sola Date: Mon, 28 Oct 2024 14:11:18 -0700 Subject: [PATCH 1/3] fix enum type assertion with python versions less than 3.12 Signed-off-by: Daniel Sola --- flytekit/core/type_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 72a59aa82e..8a08c7a2cd 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -937,7 +937,7 @@ def assert_type(self, t: Type[enum.Enum], v: T): return val = v.value if isinstance(v, enum.Enum) else v - if val not in t: + if val not in [t_item.value for t_item in t]: raise TypeTransformerFailedError(f"Value {v} is not in Enum {t}") From 05d7d9778b4dfca61e9b950e9f0759d58e9846eb Mon Sep 17 00:00:00 2001 From: Daniel Sola Date: Mon, 28 Oct 2024 14:22:21 -0700 Subject: [PATCH 2/3] add enum test Signed-off-by: Daniel Sola --- tests/flytekit/unit/core/test_enum_type.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 tests/flytekit/unit/core/test_enum_type.py diff --git a/tests/flytekit/unit/core/test_enum_type.py b/tests/flytekit/unit/core/test_enum_type.py new file mode 100644 index 0000000000..0f6881c8a3 --- /dev/null +++ b/tests/flytekit/unit/core/test_enum_type.py @@ -0,0 +1,22 @@ +from enum import Enum + +from flytekit import task + + +def test_dynamic_local(): + class Color(Enum): + RED = 'red' + GREEN = 'green' + BLUE = 'blue' + + @task + def my_task(c: Color) -> Color: + print(c) + return c + + @workflow + def wf(c: Color) -> Color: + return my_task(c=c) + + res = wf(c=Color.RED) + assert res == Color.RED From 68d08ccd9a0f3cfc0a99f4f46419cf46a3c01451 Mon Sep 17 00:00:00 2001 From: Daniel Sola Date: Mon, 28 Oct 2024 14:23:30 -0700 Subject: [PATCH 3/3] add enum test Signed-off-by: Daniel Sola --- tests/flytekit/unit/core/test_enum_type.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_enum_type.py b/tests/flytekit/unit/core/test_enum_type.py index 0f6881c8a3..2c41549ded 100644 --- a/tests/flytekit/unit/core/test_enum_type.py +++ b/tests/flytekit/unit/core/test_enum_type.py @@ -1,6 +1,6 @@ from enum import Enum -from flytekit import task +from flytekit import task, workflow def test_dynamic_local():