From eeedc0c56462f1b1dca47637b34477a027af3308 Mon Sep 17 00:00:00 2001
From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com>
Date: Fri, 13 Feb 2026 17:11:50 -0800
Subject: [PATCH 1/2] added support for enum types

Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com>
---
 flytekit/core/type_engine.py                  |  4 ++
 .../unit/core/test_enum_in_pydantic.py        | 45 +++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 tests/flytekit/unit/core/test_enum_in_pydantic.py

diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
index 3e436f0814..0e20e379f9 100644
--- a/flytekit/core/type_engine.py
+++ b/flytekit/core/type_engine.py
@@ -1233,6 +1233,10 @@ def generate_attribute_list_from_dataclass_json_mixin(
             defs = schema.get("$defs", schema.get("definitions", {}))
             if ref_name in defs:
                 ref_schema = defs[ref_name].copy()
+                # Check if the $ref points to an enum definition (no properties)
+                if ref_schema.get("enum"):
+                    attribute_list.append((property_key, str))
+                    continue
                 # Include $defs so nested models can resolve their own $refs
                 if "$defs" not in ref_schema and defs:
                     ref_schema["$defs"] = defs
diff --git a/tests/flytekit/unit/core/test_enum_in_pydantic.py b/tests/flytekit/unit/core/test_enum_in_pydantic.py
new file mode 100644
index 0000000000..5ad49331d7
--- /dev/null
+++ b/tests/flytekit/unit/core/test_enum_in_pydantic.py
@@ -0,0 +1,45 @@
+import dataclasses
+from enum import Enum
+
+from pydantic import BaseModel
+
+from flytekit.core.context_manager import FlyteContext
+from flytekit.core.type_engine import TypeEngine
+
+
+class Status(str, Enum):
+    PENDING = "pending"
+    APPROVED = "approved"
+    REJECTED = "rejected"
+
+
+class Job(BaseModel):
+    name: str
+    status: Status
+
+
+def test_pydantic_model_with_enum_ref():
+    """Test that a Pydantic model with an enum field (which produces a $ref in
+    the JSON schema) can be round-tripped through the type engine and that
+    guess_python_type reconstructs a valid dataclass."""
+    ctx = FlyteContext.current_context()
+    input = Job(name="test-job", status=Status.PENDING)
+
+    lt = TypeEngine.to_literal_type(Job)
+    lv = TypeEngine.to_literal(ctx, input, Job, lt)
+
+    assert lt
+    assert lv
+
+    # Roundtrip via the real Pydantic model
+    pv = TypeEngine.to_python_value(ctx, lv, Job)
+    assert pv == input
+
+    # Guess python type from the schema (simulates pyflyte run behaviour)
+    guessed = TypeEngine.guess_python_type(lt)
+    assert dataclasses.is_dataclass(guessed)
+
+    # The enum field should be reconstructed as str (enum $ref resolved)
+    v = guessed(name="test-job", status="pending")
+    assert v.name == "test-job"
+    assert v.status == "pending"
\ No newline at end of file

From 83826a5d936920688df3c07c172f508af1ddc434 Mon Sep 17 00:00:00 2001
From: "M. Adil Fayyaz" <62440954+AdilFayyaz@users.noreply.github.com>
Date: Fri, 13 Feb 2026 18:30:19 -0800
Subject: [PATCH 2/2] fix: lint

Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com>
---
 tests/flytekit/unit/core/test_enum_in_pydantic.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/flytekit/unit/core/test_enum_in_pydantic.py b/tests/flytekit/unit/core/test_enum_in_pydantic.py
index 5ad49331d7..76f22b79e2 100644
--- a/tests/flytekit/unit/core/test_enum_in_pydantic.py
+++ b/tests/flytekit/unit/core/test_enum_in_pydantic.py
@@ -42,4 +42,4 @@ def test_pydantic_model_with_enum_ref():
     # The enum field should be reconstructed as str (enum $ref resolved)
     v = guessed(name="test-job", status="pending")
     assert v.name == "test-job"
-    assert v.status == "pending"
\ No newline at end of file
+    assert v.status == "pending"