flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/types/_type_engine.py

@@ -16,13 +16,13 @@ import typing
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from functools import lru_cache
-from types import GenericAlias
+from types import GenericAlias, NoneType
 from typing import Any, Dict, NamedTuple, Optional, Type, cast
 
 import msgpack
-from flyteidl.core import interface_pb2, literals_pb2, types_pb2
-from flyteidl.core.literals_pb2 import Binary, Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Union, Void
-from flyteidl.core.types_pb2 import LiteralType, SimpleType, TypeAnnotation, TypeStructure, UnionType
+from flyteidl2.core import interface_pb2, literals_pb2, types_pb2
+from flyteidl2.core.literals_pb2 import Binary, Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Union, Void
+from flyteidl2.core.types_pb2 import LiteralType, SimpleType, TypeAnnotation, TypeStructure, UnionType
 from fsspec.asyn import _run_coros_in_chunks  # pylint: disable=W0212
 from google.protobuf import json_format as _json_format
 from google.protobuf import struct_pb2
@@ -35,13 +35,14 @@ from mashumaro.jsonschema.models import Context, JSONSchema
 from mashumaro.jsonschema.plugins import BasePlugin
 from mashumaro.jsonschema.schema import Instance
 from mashumaro.mixins.json import DataClassJSONMixin
+from pydantic import BaseModel
 from typing_extensions import Annotated, get_args, get_origin
 
 import flyte.storage as storage
-from flyte._datastructures import NativeInterface
-from flyte._hash import HashMethod
 from flyte._logging import logger
 from flyte._utils.helpers import load_proto_from_file
+from flyte.errors import RestrictedTypeError
+from flyte.models import NativeInterface
 
 from ._utils import literal_types_match
 
@@ -306,6 +307,9 @@ class SimpleTransformer(TypeTransformer[T]):
         expected_python_type = get_underlying_type(expected_python_type)
 
         if expected_python_type is not self._type:
+            if expected_python_type is None and issubclass(self._type, NoneType):
+                # If the expected type is NoneType, we can return None
+                return None
             raise TypeTransformerFailedError(
                 f"Cannot convert to type {expected_python_type}, only {self._type} is supported"
             )
@@ -328,10 +332,6 @@ class SimpleTransformer(TypeTransformer[T]):
         raise ValueError(f"Transformer {self} cannot reverse {literal_type}")
 
 
-class RestrictedTypeError(Exception):
-    pass
-
-
 class RestrictedTypeTransformer(TypeTransformer[T], ABC):
     """
     Types registered with the RestrictedTypeTransformer are not allowed to be converted to and from literals.
@@ -352,6 +352,66 @@ class RestrictedTypeTransformer(TypeTransformer[T], ABC):
         raise RestrictedTypeError(f"Transformer for type {self.python_type} is restricted currently")
 
 
+class PydanticTransformer(TypeTransformer[BaseModel]):
+    def __init__(self):
+        super().__init__("Pydantic Transformer", BaseModel, enable_type_assertions=False)
+
+    def get_literal_type(self, t: Type[BaseModel]) -> LiteralType:
+        schema = t.model_json_schema()
+
+        meta_struct = struct_pb2.Struct()
+        meta_struct.update(
+            {
+                CACHE_KEY_METADATA: {
+                    SERIALIZATION_FORMAT: MESSAGEPACK,
+                }
+            }
+        )
+
+        # The type engine used to publish a type structure for attribute access. As of v2, this is no longer needed.
+        return LiteralType(
+            simple=SimpleType.STRUCT,
+            metadata=schema,
+            annotation=TypeAnnotation(annotations=meta_struct),
+        )
+
+    async def to_literal(
+        self,
+        python_val: BaseModel,
+        python_type: Type[BaseModel],
+        expected: LiteralType,
+    ) -> Literal:
+        json_str = python_val.model_dump_json()
+        dict_obj = json.loads(json_str)
+        msgpack_bytes = msgpack.dumps(dict_obj)
+        return Literal(scalar=Scalar(binary=Binary(value=msgpack_bytes, tag=MESSAGEPACK)))
+
+    def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[BaseModel]) -> BaseModel:
+        if binary_idl_object.tag == MESSAGEPACK:
+            dict_obj = msgpack.loads(binary_idl_object.value, strict_map_key=False)
+            json_str = json.dumps(dict_obj)
+            python_val = expected_python_type.model_validate_json(
+                json_data=json_str, strict=False, context={"deserialize": True}
+            )
+            return python_val
+        else:
+            raise TypeTransformerFailedError(f"Unsupported binary format: `{binary_idl_object.tag}`")
+
+    async def to_python_value(self, lv: Literal, expected_python_type: Type[BaseModel]) -> BaseModel:
+        """
+        There are two kinds of literal values to handle:
+        1. Protobuf Structs (from the UI)
+        2. Binary scalars (from other sources)
+        We need to account for both cases accordingly.
+        """
+        if lv and lv.HasField("scalar") and lv.scalar.HasField("binary"):
+            return self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
+
+        json_str = _json_format.MessageToJson(lv.scalar.generic)
+        python_val = expected_python_type.model_validate_json(json_str, strict=False, context={"deserialize": True})
+        return python_val
+
+
 class PydanticSchemaPlugin(BasePlugin):
     """This allows us to generate proper schemas for Pydantic models."""
 
@@ -392,35 +452,6 @@ class DataclassTransformer(TypeTransformer[object]):
     2. Deserialization: The dataclass transformer converts the MessagePack Bytes back to a dataclass.
         (1) Convert MessagePack Bytes to a dataclass using mashumaro.
         (2) Handle dataclass attributes to ensure they are of the correct types.
-
-    TODO: Update the example using mashumaro instead of the older library
-
-    Example
-
-    .. code-block:: python
-
-        @dataclass
-        class Test:
-           a: int
-           b: str
-
-        t = Test(a=10,b="e")
-        JSONSchema().dump(t.schema())
-
-    Output will look like
-
-    .. code-block:: json
-
-        {'$schema': 'http://json-schema.org/draft-07/schema#',
-         'definitions': {'TestSchema': {'properties': {'a': {'title': 'a',
-             'type': 'number',
-             'format': 'integer'},
-            'b': {'title': 'b', 'type': 'string'}},
-           'type': 'object',
-           'additionalProperties': False}},
-         '$ref': '#/definitions/TestSchema'}
-
-
     """
 
     def __init__(self) -> None:
@@ -450,7 +481,6 @@ class DataclassTransformer(TypeTransformer[object]):
         # However, FooSchema is created by flytekit and it's not equal to the user-defined dataclass (Foo).
         # Therefore, we should iterate all attributes in the dataclass and check the type of value in dataclass
         #   matches the expected_type.
-
         expected_fields_dict = {}
 
         for f in dataclasses.fields(expected_type):
@@ -544,28 +574,7 @@ class DataclassTransformer(TypeTransformer[object]):
                 f"Possibly remove `DataClassJsonMixin` and `dataclass_json` decorator from dataclass declaration"
             )
 
-        # Recursively construct the dataclass_type which contains the literal type of each field
-        literal_type = {}
-
-        hints = typing.get_type_hints(t)
-        # Get the type of each field from dataclass
-        for field in t.__dataclass_fields__.values():  # type: ignore
-            try:
-                name = field.name
-                python_type = hints.get(name, field.type)
-                literal_type[name] = TypeEngine.to_literal_type(python_type)
-            except Exception as e:
-                logger.warning(
-                    "Field {} of type {} cannot be converted to a literal type. Error: {}".format(
-                        field.name, field.type, e
-                    )
-                )
-
-        # This is for attribute access in FlytePropeller.
-        ts = TypeStructure(tag="", dataclass_type=literal_type)
-        from google.protobuf.struct_pb2 import Struct
-
-        meta_struct = Struct()
+        meta_struct = struct_pb2.Struct()
         meta_struct.update(
             {
                 CACHE_KEY_METADATA: {
@@ -573,10 +582,11 @@ class DataclassTransformer(TypeTransformer[object]):
                 }
             }
         )
+
+        # The type engine used to publish the type `structure` for attribute access. As of v2, this is no longer needed.
         return types_pb2.LiteralType(
             simple=types_pb2.SimpleType.STRUCT,
             metadata=schema,
-            structure=ts,
             annotation=TypeAnnotation(annotations=meta_struct),
         )
 
@@ -628,7 +638,7 @@ class DataclassTransformer(TypeTransformer[object]):
                 field.type = self._get_origin_type_in_annotation(cast(type, field.type))
         return python_type
 
-    async def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[T]) -> T:
+    def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[T]) -> T:
         if binary_idl_object.tag == MESSAGEPACK:
             if issubclass(expected_python_type, DataClassJSONMixin):
                 dict_obj = msgpack.loads(binary_idl_object.value, strict_map_key=False)
@@ -653,9 +663,10 @@ class DataclassTransformer(TypeTransformer[object]):
                 "user defined datatypes in Flytekit"
             )
 
-        if lv.scalar and lv.scalar.binary:
-            return await self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
+        if lv.HasField("scalar") and lv.scalar.HasField("binary"):
+            return self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
 
+        # todo: revisit this, it should always be a binary in v2.
         json_str = _json_format.MessageToJson(lv.scalar.generic)
 
         # The `from_json` function is provided from mashumaro's `DataClassJSONMixin`.
@@ -772,6 +783,13 @@ class EnumTransformer(TypeTransformer[enum.Enum]):
         return LiteralType(enum_type=types_pb2.EnumType(values=values))
 
     async def to_literal(self, python_val: enum.Enum, python_type: Type[T], expected: LiteralType) -> Literal:
+        if isinstance(python_val, str):
+            # this is the case when python Literals are used as enums
+            if python_val not in expected.enum_type.values:
+                raise TypeTransformerFailedError(
+                    f"Value {python_val} is not valid value, expected - {expected.enum_type.values}"
+                )
+            return Literal(scalar=Scalar(primitive=Primitive(string_value=python_val)))  # type: ignore
         if type(python_val).__class__ != enum.EnumMeta:
             raise TypeTransformerFailedError("Expected an enum")
         if type(python_val.value) is not str:
@@ -782,6 +800,12 @@ class EnumTransformer(TypeTransformer[enum.Enum]):
     async def to_python_value(self, lv: Literal, expected_python_type: Type[T]) -> T:
         if lv.HasField("scalar") and lv.scalar.HasField("binary"):
             return self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
+        from flyte._interface import LITERAL_ENUM
+
+        if expected_python_type.__name__ is LITERAL_ENUM:
+            # This is the case when python Literal types are used as enums. The class name is always LiteralEnum an
+            # hardcoded in flyte.models
+            return lv.scalar.primitive.string_value
         return expected_python_type(lv.scalar.primitive.string_value)  # type: ignore
 
     def guess_python_type(self, literal_type: LiteralType) -> Type[enum.Enum]:
@@ -800,7 +824,37 @@ 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]] = []
-    for property_key, property_val in schema["properties"].items():
+    nested_types: typing.Dict[str, type] = {}  # Track nested model types for conversion
+
+    # Use 'required' field to preserve property order, as protobuf Struct doesn't preserve dict order
+    properties = schema["properties"]
+    property_order = schema.get("required", list(properties.keys()))
+
+    for property_key in property_order:
+        property_val = properties[property_key]
+        # Handle $ref for nested Pydantic models
+        if property_val.get("$ref"):
+            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()
+                # Include $defs so nested models can resolve their own $refs
+                if "$defs" not in ref_schema and defs:
+                    ref_schema["$defs"] = defs
+                nested_class: type = convert_mashumaro_json_schema_to_python_class(ref_schema, ref_name)
+                attribute_list.append(
+                    (
+                        property_key,
+                        typing.cast(GenericAlias, nested_class),
+                    )
+                )
+                # Track this as a nested type that needs dict-to-object conversion
+                nested_types[property_key] = nested_class
+            continue
+
         if property_val.get("anyOf"):
             property_type = property_val["anyOf"][0]["type"]
         elif property_val.get("enum"):
@@ -838,14 +892,14 @@ def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name:
                         )
                     )
                     continue
+                nested_class = convert_mashumaro_json_schema_to_python_class(sub_schemea, sub_schemea_name)
                 attribute_list.append(
                     (
                         property_key,
-                        typing.cast(
-                            GenericAlias, convert_mashumaro_json_schema_to_python_class(sub_schemea, sub_schemea_name)
-                        ),
+                        typing.cast(GenericAlias, nested_class),
                     )
                 )
+                nested_types[property_key] = nested_class
             elif property_val.get("additionalProperties"):
                 # For typing.Dict type
                 elem_type = _get_element_type(property_val["additionalProperties"])
@@ -876,14 +930,14 @@ def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name:
                         )
                     )
                     continue
+                nested_class = convert_mashumaro_json_schema_to_python_class(property_val, sub_schemea_name)
                 attribute_list.append(
                     (
                         property_key,
-                        typing.cast(
-                            GenericAlias, convert_mashumaro_json_schema_to_python_class(property_val, sub_schemea_name)
-                        ),
+                        typing.cast(GenericAlias, nested_class),
                     )
                 )
+                nested_types[property_key] = nested_class
             else:
                 # For untyped dict
                 attribute_list.append((property_key, dict))  # type: ignore
@@ -892,7 +946,7 @@ def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name:
         # Handle int, float, bool or str
         else:
             attribute_list.append([property_key, _get_element_type(property_val)])  # type: ignore
-    return attribute_list
+    return attribute_list, nested_types
 
 
 class TypeEngine(typing.Generic[T]):
@@ -971,11 +1025,10 @@ class TypeEngine(typing.Generic[T]):
                 return cls._REGISTRY[python_type.__origin__]
 
         # Handling UnionType specially - PEP 604
-        if sys.version_info >= (3, 10):
-            import types
+        import types
 
-            if isinstance(python_type, types.UnionType):
-                return cls._REGISTRY[types.UnionType]
+        if isinstance(python_type, types.UnionType):
+            return cls._REGISTRY[types.UnionType]
 
         if python_type in cls._REGISTRY:
             return cls._REGISTRY[python_type]
@@ -1009,7 +1062,7 @@ class TypeEngine(typing.Generic[T]):
             return cls._DATACLASS_TRANSFORMER
 
         display_pickle_warning(str(python_type))
-        from flyte.io.pickle.transformer import FlytePickleTransformer
+        from flyte.types._pickle import FlytePickleTransformer
 
         return FlytePickleTransformer()
 
@@ -1022,10 +1075,10 @@ class TypeEngine(typing.Generic[T]):
             # Avoid a race condition where concurrent threads may exit lazy_import_transformers before the transformers
             # have been imported. This could be implemented without a lock if you assume python assignments are atomic
             # and re-registering transformers is acceptable, but I decided to play it safe.
-            from flyte.io.structured_dataset import lazy_import_structured_dataset_handler
+            from flyte.io._dataframe import lazy_import_dataframe_handler
 
             # todo: bring in extras transformers (pytorch, etc.)
-            lazy_import_structured_dataset_handler()
+            lazy_import_dataframe_handler()
 
     @classmethod
     def to_literal_type(cls, python_type: Type[T]) -> LiteralType:
@@ -1053,23 +1106,6 @@ class TypeEngine(typing.Generic[T]):
         ):
             raise TypeTransformerFailedError(f"Python value cannot be None, expected {python_type}/{expected}")
 
-    @classmethod
-    def calculate_hash(cls, python_val: typing.Any, python_type: Type[T]) -> Optional[str]:
-        # In case the value is an annotated type we inspect the annotations and look for hash-related annotations.
-        hsh = None
-        if is_annotated(python_type):
-            # We are now dealing with one of two cases:
-            # 1. The annotated type is a `HashMethod`, which indicates that we should produce the hash using
-            #    the method indicated in the annotation.
-            # 2. The annotated type is being used for a different purpose other than calculating hash values,
-            #    in which case we should just continue.
-            for annotation in get_args(python_type)[1:]:
-                if not isinstance(annotation, HashMethod):
-                    continue
-                hsh = annotation.calculate(python_val)
-                break
-        return hsh
-
     @classmethod
     async def to_literal(
         cls, python_val: typing.Any, python_type: Type[T], expected: types_pb2.LiteralType
@@ -1082,8 +1118,6 @@ class TypeEngine(typing.Generic[T]):
         lv = await transformer.to_literal(python_val, python_type, expected)
 
         modify_literal_uris(lv)
-        calculated_hash = cls.calculate_hash(python_val, python_type) or ""
-        lv.hash = calculated_hash
         return lv
 
     @classmethod
@@ -1161,20 +1195,26 @@ class TypeEngine(typing.Generic[T]):
                 f"Received more input values {len(lm.literals)}"
                 f" than allowed by the input spec {len(python_interface_inputs)}"
             )
+        # Create tasks for converting each kwarg
+        tasks = {}
+        for k in lm.literals:
+            tasks[k] = asyncio.create_task(TypeEngine.to_python_value(lm.literals[k], python_interface_inputs[k]))
+
+        # Gather all tasks, returning exceptions instead of raising them
+        results = await asyncio.gather(*tasks.values(), return_exceptions=True)
+
+        # Check for exceptions and raise with specific kwarg name
         kwargs = {}
-        try:
-            for i, k in enumerate(lm.literals):
-                kwargs[k] = asyncio.create_task(TypeEngine.to_python_value(lm.literals[k], python_interface_inputs[k]))
-            await asyncio.gather(*kwargs.values())
-        except Exception as e:
-            raise TypeTransformerFailedError(
-                f"Error converting input:\n"
-                f"Literal value: {lm.literals[k]}\n"
-                f"Expected Python type: {python_interface_inputs[k]}\n"
-                f"Exception: {e}"
-            )
+        for (key, task), result in zip(tasks.items(), results):
+            if isinstance(result, Exception):
+                raise TypeTransformerFailedError(
+                    f"Error converting input '{key}':\n"
+                    f"Literal value: {lm.literals[key]}\n"
+                    f"Expected Python type: {python_interface_inputs[key]}\n"
+                    f"Exception: {result}"
+                ) from result
+            kwargs[key] = result
 
-        kwargs = {k: v.result() for k, v in kwargs.items() if v is not None}
         return kwargs
 
     @classmethod
@@ -1208,7 +1248,12 @@ class TypeEngine(typing.Generic[T]):
                 python_type = type_hints.get(k, type(d[k]))
                 e: BaseException = literal_map[k].exception()  # type: ignore
                 if isinstance(e, TypeError):
-                    raise TypeError(f"Error converting: {type(v)}, {python_type}, received_value {v}")
+                    raise TypeError(
+                        f"Type conversion failed for variable '{k}'.\n"
+                        f"Expected type: {python_type}\n"
+                        f"Actual type: {type(d[k])}\n"
+                        f"Value received: {d[k]!r}"
+                    ) from e
                 else:
                     raise e
             literal_map[k] = v.result()
@@ -1243,7 +1288,8 @@ class TypeEngine(typing.Generic[T]):
             try:
                 return transformer.guess_python_type(flyte_type)
             except ValueError:
-                logger.debug(f"Skipping transformer {transformer.name} for {flyte_type}")
+                # Skipping transformer
+                continue
 
         # Because the dataclass transformer is handled explicitly in the get_transformer code, we have to handle it
         # separately here too.
@@ -1366,6 +1412,19 @@ def _type_essence(x: types_pb2.LiteralType) -> types_pb2.LiteralType:
 
 
 def _are_types_castable(upstream: types_pb2.LiteralType, downstream: types_pb2.LiteralType) -> bool:
+    if upstream.union_type is not None:
+        # for each upstream variant, there must be a compatible type downstream
+        for v in upstream.union_type.variants:
+            if not _are_types_castable(v, downstream):
+                return False
+        return True
+
+    if downstream.union_type is not None:
+        # there must be a compatible downstream type
+        for v in downstream.union_type.variants:
+            if _are_types_castable(upstream, v):
+                return True
+
     if upstream.HasField("collection_type"):
         if not downstream.HasField("collection_type"):
             return False
@@ -1411,7 +1470,7 @@ def _are_types_castable(upstream: types_pb2.LiteralType, downstream: types_pb2.L
 
         return True
 
-    if upstream.HasField("union_type"):
+    if upstream.HasField("union_type") and upstream.union_type is not None:
         # for each upstream variant, there must be a compatible type downstream
         for v in upstream.union_type.variants:
             if not _are_types_castable(v, downstream):
@@ -1505,6 +1564,18 @@ class UnionTransformer(TypeTransformer[T]):
         self, python_val: T, python_type: Type[T], expected: types_pb2.LiteralType
     ) -> literals_pb2.Literal:
         python_type = get_underlying_type(python_type)
+        inferred_type = type(python_val)
+        subtypes = get_args(python_type)
+
+        if inferred_type in subtypes:
+            # If the Python value's type matches one of the types in the Union,
+            # always use the transformer associated with that specific type.
+            transformer = TypeEngine.get_transformer(inferred_type)
+            res = await transformer.to_literal(
+                python_val, inferred_type, expected.union_type.variants[subtypes.index(inferred_type)]
+            )
+            res_type = _add_tag_to_type(transformer.get_literal_type(inferred_type), transformer.name)
+            return Literal(scalar=Scalar(union=Union(value=res, type=res_type)))
 
         potential_types = []
         found_res = False
@@ -1512,14 +1583,14 @@ class UnionTransformer(TypeTransformer[T]):
         res = None
         res_type = None
         t = None
-        for i in range(len(get_args(python_type))):
+        for i in range(len(subtypes)):
             try:
-                t = get_args(python_type)[i]
+                t = subtypes[i]
                 trans: TypeTransformer[T] = TypeEngine.get_transformer(t)
                 attempt = trans.to_literal(python_val, t, expected.union_type.variants[i])
                 res = await attempt
                 if found_res:
-                    logger.debug(f"Current type {get_args(python_type)[i]} old res {res_type}")
+                    logger.debug(f"Current type {subtypes[i]} old res {res_type}")
                     is_ambiguous = True
                 res_type = _add_tag_to_type(trans.get_literal_type(t), trans.name)
                 found_res = True
@@ -1658,7 +1729,7 @@ class DictTransformer(TypeTransformer[dict]):
         Converts a Python dictionary to a Flyte-specific ``Literal`` using MessagePack encoding.
         Falls back to Pickle if encoding fails and `allow_pickle` is True.
         """
-        from flyte.io.pickle.transformer import FlytePickle
+        from flyte.types._pickle import FlytePickle
 
         try:
             # Handle dictionaries with non-string keys (e.g., Dict[int, Type])
@@ -1724,7 +1795,6 @@ class DictTransformer(TypeTransformer[dict]):
         for k, v in python_val.items():
             if type(k) is not str:
                 raise ValueError("Flyte MapType expects all keys to be strings")
-            # TODO: log a warning for Annotated objects that contain HashMethod
 
             _, v_type = self.extract_types(python_type)
             lit_map[k] = TypeEngine.to_literal(v, cast(type, v_type), expected.map_value_type)
@@ -1764,7 +1834,7 @@ class DictTransformer(TypeTransformer[dict]):
         # pr: han-ru is this part still necessary?
         if lv and lv.HasField("scalar") and lv.scalar.HasField("generic"):
             if lv.metadata and lv.metadata.get("format", None) == "pickle":
-                from flyte.io.pickle.transformer import FlytePickle
+                from flyte.types._pickle import FlytePickle
 
                 uri = json.loads(_json_format.MessageToJson(lv.scalar.generic)).get("pickle_file")
                 return await FlytePickle.from_pickle(uri)
@@ -1834,8 +1904,26 @@ def convert_mashumaro_json_schema_to_python_class(schema: dict, schema_name: typ
     :param schema_name: dataclass name of return type
     """
 
-    attribute_list = generate_attribute_list_from_dataclass_json_mixin(schema, schema_name)
-    return dataclasses.make_dataclass(schema_name, attribute_list)
+    attribute_list, nested_types = generate_attribute_list_from_dataclass_json_mixin(schema, schema_name)
+    cls = dataclasses.make_dataclass(schema_name, attribute_list)
+
+    # Wrap __init__ to convert dict inputs to nested types
+    if nested_types:
+        # Store the original __init__ from the class's __dict__ to avoid mypy error
+        original_init = cls.__dict__["__init__"]
+
+        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():
+                if field_name in kwargs:
+                    value = kwargs[field_name]
+                    if isinstance(value, dict):
+                        kwargs[field_name] = field_type(**value)
+            original_init(self, *args, **kwargs)
+
+        cls.__init__ = __init__  # type: ignore[method-assign, misc]
+
+    return cls
 
 
 def _get_element_type(element_property: typing.Dict[str, str]) -> Type:
@@ -1871,7 +1959,6 @@ def _get_element_type(element_property: typing.Dict[str, str]) -> Type:
     return str
 
 
-# pr: han-ru is this still needed?
 def dataclass_from_dict(cls: type, src: typing.Dict[str, typing.Any]) -> typing.Any:
     """
     Utility function to construct a dataclass object from dict
@@ -1951,7 +2038,7 @@ def _handle_flyte_console_float_input_to_int(lv: Literal) -> int:
 
 def _check_and_convert_void(lv: Literal) -> None:
     if not lv.scalar.HasField("none_type"):
-        raise TypeTransformerFailedError(f"Cannot convert literal {lv} to None")
+        raise TypeTransformerFailedError(f"Cannot convert literal '{lv}' to None")
     return None
 
 
@@ -1992,7 +2079,9 @@ DatetimeTransformer = SimpleTransformer(
     datetime.datetime,
     types_pb2.LiteralType(simple=types_pb2.SimpleType.DATETIME),
     lambda x: Literal(scalar=Scalar(primitive=Primitive(datetime=x))),
-    lambda x: x.scalar.primitive.datetime if x.scalar.primitive.HasField("datetime") else None,
+    lambda x: x.scalar.primitive.datetime.ToDatetime().replace(tzinfo=datetime.timezone.utc)
+    if x.scalar.primitive.HasField("datetime")
+    else None,
 )
 
 TimedeltaTransformer = SimpleTransformer(
@@ -2000,7 +2089,7 @@ TimedeltaTransformer = SimpleTransformer(
     datetime.timedelta,
     types_pb2.LiteralType(simple=types_pb2.SimpleType.DURATION),
     lambda x: Literal(scalar=Scalar(primitive=Primitive(duration=x))),
-    lambda x: x.scalar.primitive.duration if x.scalar.primitive.HasField("duration") else None,
+    lambda x: x.scalar.primitive.duration.ToTimedelta() if x.scalar.primitive.HasField("duration") else None,
 )
 
 DateTransformer = SimpleTransformer(
@@ -2010,7 +2099,9 @@ DateTransformer = SimpleTransformer(
     lambda x: Literal(
         scalar=Scalar(primitive=Primitive(datetime=datetime.datetime.combine(x, datetime.time.min)))
     ),  # convert datetime to date
-    lambda x: x.scalar.primitive.datetime.date() if x.scalar.primitive.HasField("datetime") else None,
+    lambda x: x.scalar.primitive.datetime.ToDatetime().replace(tzinfo=datetime.timezone.utc).date()
+    if x.scalar.primitive.HasField("datetime")
+    else None,
 )
 
 NoneTransformer = SimpleTransformer(
@@ -2034,10 +2125,20 @@ def _register_default_type_transformers():
     TypeEngine.register(BoolTransformer)
     TypeEngine.register(NoneTransformer, [None])
     TypeEngine.register(ListTransformer())
-    TypeEngine.register(UnionTransformer(), [UnionType])
+
+    if sys.version_info < (3, 14):
+        TypeEngine.register(UnionTransformer(), [UnionType])
+    else:
+        # In Python 3.14+, types.UnionType and typing.Union are the same object.
+        # UnionTransformer's python_type is already typing.Union, so only add UnionType
+        # as an additional type if it's different from typing.Union.
+        union_transformer = UnionTransformer()
+        additional_union_types = [] if UnionType is union_transformer.python_type else [UnionType]
+        TypeEngine.register(union_transformer, additional_union_types)
     TypeEngine.register(DictTransformer())
     TypeEngine.register(EnumTransformer())
     TypeEngine.register(ProtobufTransformer())
+    TypeEngine.register(PydanticTransformer())
 
     # inner type is. Also unsupported are typing's Tuples. Even though you can look inside them, Flyte's type system
     # doesn't support these currently.

flyte/types/_utils.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import importlib
 import typing
 
-from flyteidl.core.types_pb2 import EnumType, LiteralType, UnionType
+from flyteidl2.core.types_pb2 import EnumType, LiteralType, UnionType
 
 T = typing.TypeVar("T")