Tuple fix (#684)

The pydantic transformer to support non-native Flyte types was erroring
on `complex_typeddict_workflow` and `optional_fields_workflow`. All the
complex workflow example now run.

---------

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

Author: wild-endeavor

src/flyte/types/_tuple_dict.py

@@ -256,6 +256,8 @@ def _convert_value_from_pydantic(self, value: Any, expected_type: Optional[Type]
             return [self._convert_value_from_pydantic(v, args[0]) for v in value]
         elif isinstance(value, dict) and origin is dict and len(args) >= 2:
             return {k: self._convert_value_from_pydantic(v, args[1]) for k, v in value.items()}
+        elif isinstance(value, (list, tuple)) and (origin is tuple or origin is Tuple) and args:
+            return tuple(self._convert_value_from_pydantic(v, t) for v, t in zip(value, args))
 
         return value
 
@@ -611,8 +613,13 @@ def _create_wrapper_model(
             return _model_cache[t]
 
         # Get field names and types from the TypedDict
-        annotations = getattr(t, "__annotations__", {})
-        required_keys: frozenset[str] = getattr(t, "__required_keys__", frozenset())
+        # Use get_type_hints(include_extras=True) to preserve NotRequired/Required wrappers,
+        # which we need to detect optional fields. We can't rely on __required_keys__ because
+        # it is wrong when `from __future__ import annotations` is used (all fields appear required).
+        try:
+            annotations_with_extras = typing.get_type_hints(t, include_extras=True)
+        except Exception:
+            annotations_with_extras = getattr(t, "__annotations__", {})
 
         field_definitions: Dict[str, Any] = {}
 
@@ -621,11 +628,18 @@ def _create_wrapper_model(
 
         model_name = f"{self._WRAPPER_PREFIX}{t.__name__}"
 
-        # Create a placeholder in the cache before processing fields to handle self-referential types
-        # We'll update it after the model is created
-        _model_cache[t] = None  # type: ignore
+        # Use the model name as a placeholder in the cache before processing fields.
+        # For self-referential types (e.g. TreeNode with children: List[TreeNode]),
+        # the recursive call to _create_wrapper_model will return this string, which
+        # becomes a forward reference (e.g. List["TypedDictWrapper_TreeNode"]) that
+        # Pydantic resolves via model_rebuild() after model creation.
+        _model_cache[t] = model_name  # type: ignore
+
+        for field_name, field_type in annotations_with_extras.items():
+            # Check if the field is NotRequired before unwrapping
+            origin = get_origin(field_type)
+            is_not_required = origin is NotRequired
 
-        for field_name, field_type in annotations.items():
             # Unwrap NotRequired and Required type hints to get the inner type
             # These are only used by TypedDict to mark optional/required fields
             # and should not be passed to Pydantic
@@ -635,14 +649,19 @@ def _create_wrapper_model(
             # This is necessary because isinstance() doesn't work with TypedDict on Python < 3.12
             pydantic_type = self._convert_field_type_for_pydantic(inner_type, _model_cache)
 
-            if field_name in required_keys:
-                field_definitions[field_name] = (pydantic_type, ...)
-            else:
+            if is_not_required:
                 # Optional fields get a default of None
                 field_definitions[field_name] = (typing.Optional[pydantic_type], None)
+            else:
+                field_definitions[field_name] = (pydantic_type, ...)
 
         model = create_model(model_name, **field_definitions)
         _model_cache[t] = model
+
+        # Rebuild to resolve any forward references from self-referential types
+        rebuild_ns = {f"{self._WRAPPER_PREFIX}{k.__name__}": v for k, v in _model_cache.items() if isinstance(v, type)}
+        model.model_rebuild(_types_namespace=rebuild_ns)
+
         return model
 
     def _convert_field_type_for_pydantic(self, field_type: Type, _model_cache: Dict[Type, Type[BaseModel]]) -> Type:
@@ -701,7 +720,10 @@ def _unwrap_typeddict_field_type(self, field_type: Type) -> Type:
 
     def _value_to_model(self, python_val: dict, model_class: Type[BaseModel], python_type: Type) -> BaseModel:
         """Convert a TypedDict to a Pydantic model instance."""
-        annotations = getattr(python_type, "__annotations__", {})
+        try:
+            annotations = typing.get_type_hints(python_type)
+        except Exception:
+            annotations = getattr(python_type, "__annotations__", {})
 
         # Convert nested values that might be TypedDicts, dataclasses, or Pydantic models
         # to a format Pydantic can validate (dicts)
@@ -714,19 +736,24 @@ def _value_to_model(self, python_val: dict, model_class: Type[BaseModel], python
 
     def _model_to_value(self, model_instance: BaseModel, expected_type: Type) -> dict:
         """Convert a Pydantic model instance back to a TypedDict."""
-        annotations = getattr(expected_type, "__annotations__", {})
-        required_keys: frozenset[str] = getattr(expected_type, "__required_keys__", frozenset())
+        try:
+            annotations_with_extras = typing.get_type_hints(expected_type, include_extras=True)
+        except Exception:
+            annotations_with_extras = getattr(expected_type, "__annotations__", {})
         result = {}
-        for name, field_type in annotations.items():
+        for name, field_type in annotations_with_extras.items():
             if hasattr(model_instance, name):
                 value = getattr(model_instance, name)
                 # Skip NotRequired fields when value is None
                 # This ensures that optional fields not provided in the input
                 # are absent from the output dict (not set to None)
-                if name not in required_keys and value is None:
+                origin = get_origin(field_type)
+                if origin is NotRequired and value is None:
                     continue
+                # Unwrap NotRequired/Required before converting
+                inner_type = self._unwrap_typeddict_field_type(field_type)
                 # Recursively convert nested values
-                converted_value = self._convert_value_from_pydantic(value, field_type)
+                converted_value = self._convert_value_from_pydantic(value, inner_type)
                 result[name] = converted_value
         return result
 

tests/flyte/type_engine/test_pydantic_nonflyte_types.py

@@ -0,0 +1,80 @@
+"""
+Special test cases to test the pydantic non-flyte type transformer.
+
+1. Tuple containing TypedDict elements not round-tripping (elements stayed as Pydantic wrappers)
+2. NotRequired fields broken with `from __future__ import annotations` (__required_keys__ wrong)
+3. Self-referential TypedDicts failing (None placeholder used as a type)
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import pytest
+from typing_extensions import NotRequired, TypedDict
+
+from flyte.types._type_engine import TypeEngine
+
+
+class Coord(TypedDict):
+    x: float
+    y: float
+
+
+class Node(TypedDict):
+    value: str
+    children: NotRequired[List[Node]]
+
+
+class Message(TypedDict):
+    text: str
+    tags: NotRequired[List[str]]
+
+
+@pytest.mark.asyncio
+async def test_tuple_containing_typeddict_elements():
+    """Tuple elements that are TypedDicts must be dicts after round-trip, not Pydantic wrappers."""
+    pt = tuple[Coord, Coord]
+    lt = TypeEngine.to_literal_type(pt)
+    value = (Coord(x=1.0, y=2.0), Coord(x=3.0, y=4.0))
+
+    lv = await TypeEngine.to_literal(value, pt, lt)
+    result = await TypeEngine.to_python_value(lv, pt)
+
+    assert result[0]["x"] == 1.0
+    assert result[1]["y"] == 4.0
+
+
+@pytest.mark.asyncio
+async def test_notrequired_field_with_future_annotations():
+    """NotRequired fields must be optional even with `from __future__ import annotations`."""
+    pt = Message
+    lt = TypeEngine.to_literal_type(pt)
+
+    # Without optional field
+    value_without = Message(text="hello")
+    lv = await TypeEngine.to_literal(value_without, pt, lt)
+    result = await TypeEngine.to_python_value(lv, pt)
+    assert result == {"text": "hello"}
+    assert "tags" not in result
+
+    # With optional field
+    value_with = Message(text="hello", tags=["a", "b"])
+    lv = await TypeEngine.to_literal(value_with, pt, lt)
+    result = await TypeEngine.to_python_value(lv, pt)
+    assert result == {"text": "hello", "tags": ["a", "b"]}
+
+
+@pytest.mark.asyncio
+async def test_self_referential_typeddict():
+    """Self-referential TypedDicts must round-trip correctly."""
+    pt = Node
+    lt = TypeEngine.to_literal_type(pt)
+    value = Node(value="root", children=[Node(value="child1"), Node(value="child2")])
+
+    lv = await TypeEngine.to_literal(value, pt, lt)
+    result = await TypeEngine.to_python_value(lv, pt)
+
+    assert result["value"] == "root"
+    assert len(result["children"]) == 2
+    assert result["children"][0]["value"] == "child1"