fix: Handle list[Model], dict[str, Model] and Optional[Model] in nested type detection

The nested schema type detection from PR #426 did not resolve $ref for
list items, dict additionalProperties, or anyOf (Optional) fields.

This caused KeyError or incorrect type inference when using Pydantic or
dataclass models with fields like list[NestedModel], dict[str, NestedModel],
or Optional[NestedModel].

Changes:
- Resolve $ref in array items for list[NestedModel] support
- Resolve $ref in additionalProperties for dict[str, NestedModel] support
- Resolve $ref in anyOf for Optional[NestedModel] support
- Include optional properties in property_order so fields with defaults
  are not skipped during schema traversal
- Extend __init__ wrapper to convert list[dict] and dict[str, dict]
  to their respective nested types

Ref: flyteorg/flyte#6887
Signed-off-by: André Ahlert <andre@aex.partners>

Author: andreahlert

src/flyte/types/_type_engine.py

@@ -922,11 +922,19 @@ def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name:
     from flyte.io._file import File
 
     attribute_list: typing.List[typing.Tuple[Any, Any]] = []
-    nested_types: typing.Dict[str, type] = {}  # Track nested model types for conversion
+    # Track nested model types for conversion: (container_kind, nested_class)
+    # container_kind: "direct" for field: Model, "list" for field: list[Model], "dict" for field: dict[str, Model]
+    nested_types: typing.Dict[str, typing.Tuple[str, type]] = {}
 
     # Use 'required' field to preserve property order, as protobuf Struct doesn't preserve dict order
+    # Also include optional properties (those not in 'required') so they are not skipped
     properties = schema["properties"]
-    property_order = schema.get("required", list(properties.keys()))
+    required = schema.get("required", list(properties.keys()))
+    property_order = list(required)
+    for key in properties:
+        if key not in property_order:
+            property_order.append(key)
+    defs = schema.get("$defs", schema.get("definitions", {}))
 
     for property_key in property_order:
         property_val = properties[property_key]
@@ -935,8 +943,6 @@ def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name:
             ref_path = property_val["$ref"]
             # Extract the definition name from the $ref path (e.g., "#/$defs/MyNestedModel" -> "MyNestedModel")
             ref_name = ref_path.split("/")[-1]
-            # Get the referenced schema from $defs (or definitions for older schemas)
-            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)
@@ -954,18 +960,52 @@ def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name:
                     )
                 )
                 # Track this as a nested type that needs dict-to-object conversion
-                nested_types[property_key] = nested_class
+                nested_types[property_key] = ("direct", nested_class)
             continue
 
         if property_val.get("anyOf"):
-            property_type = property_val["anyOf"][0]["type"]
+            first_option = property_val["anyOf"][0]
+            if first_option.get("$ref"):
+                # Handle Optional[NestedModel] - anyOf with $ref and null
+                ref_path = first_option["$ref"]
+                ref_name = ref_path.split("/")[-1]
+                if ref_name in defs:
+                    ref_schema = defs[ref_name].copy()
+                    if ref_schema.get("enum"):
+                        attribute_list.append((property_key, typing.Optional[str]))
+                    else:
+                        if "$defs" not in ref_schema and defs:
+                            ref_schema["$defs"] = defs
+                        nested_class = convert_mashumaro_json_schema_to_python_class(ref_schema, ref_name)
+                        attribute_list.append((property_key, typing.Optional[typing.cast(GenericAlias, nested_class)]))
+                        nested_types[property_key] = ("direct", nested_class)
+                continue
+            property_type = first_option["type"]
         elif property_val.get("enum"):
             property_type = "enum"
         else:
             property_type = property_val["type"]
         # Handle list
         if property_type == "array":
-            attribute_list.append((property_key, typing.List[_get_element_type(property_val["items"])]))  # type: ignore
+            items = property_val["items"]
+            if isinstance(items, dict) and items.get("$ref"):
+                # Handle list[NestedModel]
+                ref_path = items["$ref"]
+                ref_name = ref_path.split("/")[-1]
+                if ref_name in defs:
+                    ref_schema = defs[ref_name].copy()
+                    if ref_schema.get("enum"):
+                        attribute_list.append((property_key, typing.List[str]))  # type: ignore
+                    else:
+                        if "$defs" not in ref_schema and defs:
+                            ref_schema["$defs"] = defs
+                        nested_class = convert_mashumaro_json_schema_to_python_class(ref_schema, ref_name)
+                        attribute_list.append((property_key, typing.List[typing.cast(GenericAlias, nested_class)]))  # type: ignore
+                        nested_types[property_key] = ("list", nested_class)
+                else:
+                    attribute_list.append((property_key, typing.List[_get_element_type(items)]))  # type: ignore
+            else:
+                attribute_list.append((property_key, typing.List[_get_element_type(items)]))  # type: ignore
         # Handle dataclass and dict
         elif property_type == "object":
             if property_val.get("anyOf"):
@@ -1001,11 +1041,31 @@ def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name:
                         typing.cast(GenericAlias, nested_class),
                     )
                 )
-                nested_types[property_key] = nested_class
+                nested_types[property_key] = ("direct", nested_class)
             elif property_val.get("additionalProperties"):
                 # For typing.Dict type
-                elem_type = _get_element_type(property_val["additionalProperties"])
-                attribute_list.append((property_key, typing.Dict[str, elem_type]))  # type: ignore
+                additional = property_val["additionalProperties"]
+                if isinstance(additional, dict) and additional.get("$ref"):
+                    # Handle dict[str, NestedModel]
+                    ref_path = additional["$ref"]
+                    ref_name = ref_path.split("/")[-1]
+                    if ref_name in defs:
+                        ref_schema = defs[ref_name].copy()
+                        if ref_schema.get("enum"):
+                            attribute_list.append((property_key, typing.Dict[str, str]))  # type: ignore
+                        else:
+                            if "$defs" not in ref_schema and defs:
+                                ref_schema["$defs"] = defs
+                            nested_class = convert_mashumaro_json_schema_to_python_class(ref_schema, ref_name)
+                            casted = typing.cast(GenericAlias, nested_class)
+                            attribute_list.append((property_key, typing.Dict[str, casted]))  # type: ignore
+                            nested_types[property_key] = ("dict", nested_class)
+                    else:
+                        elem_type = _get_element_type(additional)
+                        attribute_list.append((property_key, typing.Dict[str, elem_type]))  # type: ignore
+                else:
+                    elem_type = _get_element_type(additional)
+                    attribute_list.append((property_key, typing.Dict[str, elem_type]))  # type: ignore
             elif property_val.get("title"):
                 # For nested dataclass
                 sub_schemea_name = property_val["title"]
@@ -1039,7 +1099,7 @@ def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name:
                         typing.cast(GenericAlias, nested_class),
                     )
                 )
-                nested_types[property_key] = nested_class
+                nested_types[property_key] = ("direct", nested_class)
             else:
                 # For untyped dict
                 attribute_list.append((property_key, dict))  # type: ignore
@@ -2064,11 +2124,17 @@ def convert_mashumaro_json_schema_to_python_class(schema: dict, schema_name: typ
 
         def __init__(self, *args, **kwargs):  # type: ignore[misc]
             # Convert dict values to nested types before calling original __init__
-            for field_name, field_type in nested_types.items():
+            for field_name, (kind, field_type) in nested_types.items():
                 if field_name in kwargs:
                     value = kwargs[field_name]
-                    if isinstance(value, dict):
+                    if kind == "direct" and isinstance(value, dict):
                         kwargs[field_name] = field_type(**value)
+                    elif kind == "list" and isinstance(value, list):
+                        kwargs[field_name] = [field_type(**item) if isinstance(item, dict) else item for item in value]
+                    elif kind == "dict" and isinstance(value, dict):
+                        kwargs[field_name] = {
+                            k: field_type(**v) if isinstance(v, dict) else v for k, v in value.items()
+                        }
             original_init(self, *args, **kwargs)
 
         cls.__init__ = __init__  # type: ignore[method-assign, misc]