feat(DTO): Support nested union types with single composites (#4293)

Author: provinzkraut

litestar/dto/_codegen_backend.py

@@ -569,28 +569,79 @@ def _create_transfer_nested_union_type_data(
         source_value_name: str,
         assignment_target: str,
     ) -> None:
-        for inner_type in transfer_type.inner_types:
-            if isinstance(inner_type, CompositeType):
-                continue
+        def _handle_transfer_instance(simple_type_: SimpleType, conditional_: str) -> None:
+            if simple_type_.field_definition.is_none_type:
+                with self._start_block(f"{conditional_} {source_value_name} is None:"):
+                    self._add_stmt(f"{assignment_target} = {source_value_name}")
+                return
 
-            if inner_type.nested_field_info:
-                if self.is_data_field:
-                    constraint_type = inner_type.nested_field_info.model
-                    destination_type = inner_type.field_definition.annotation
-                else:
-                    constraint_type = inner_type.field_definition.annotation
-                    destination_type = inner_type.nested_field_info.model
+            field_definitions: tuple[TransferDTOFieldDefinition, ...] | None
+            if simple_type_.nested_field_info and self.is_data_field:
+                constraint_type = simple_type_.nested_field_info.model
+                destination_type = simple_type_.field_definition.annotation
+                field_definitions = simple_type_.nested_field_info.field_definitions
+            else:
+                constraint_type = simple_type_.field_definition.annotation
+                destination_type = (
+                    simple_type_.nested_field_info.model
+                    if simple_type_.nested_field_info and not self.is_data_field
+                    else simple_type_.field_definition.annotation
+                )
+                field_definitions = (
+                    simple_type_.nested_field_info.field_definitions if simple_type_.nested_field_info else None
+                )
 
-                constraint_type_name = self._add_to_fn_globals("constraint_type", constraint_type)
-                destination_type_name = self._add_to_fn_globals("destination_type", destination_type)
+            constraint_type_name = self._add_to_fn_globals("constraint_type", constraint_type)
+            destination_type_name = self._add_to_fn_globals("destination_type", destination_type)
 
-                with self._start_block(f"if isinstance({source_value_name}, {constraint_type_name}):"):
+            with self._start_block(f"{conditional_} isinstance({source_value_name}, {constraint_type_name}):"):
+                if field_definitions:
                     self._create_transfer_instance_data(
                         destination_type_name=destination_type_name,
                         destination_type_is_dict=destination_type is dict,
-                        field_definitions=inner_type.nested_field_info.field_definitions,
+                        field_definitions=field_definitions,
                         source_instance_name=source_value_name,
                         tmp_return_type_name=assignment_target,
                     )
-                    return
+                else:
+                    self._add_stmt(f"{assignment_target} = {source_value_name}")
+
+        simple_types: list[SimpleType] = []
+        non_simple_types: list[CompositeType] = []
+        for inner_type in transfer_type.inner_types:
+            if isinstance(inner_type, SimpleType):
+                simple_types.append(inner_type)
+            else:
+                non_simple_types.append(inner_type)
+
+        if len(non_simple_types) > 1:
+            # we've got something like 'Union[list[str], dict[str, int]]. Since checking against these goes beyond the
+            # scope of simple 'isinstance' or 'type' checks, we cannot generate code that handles these correctly.
+            # so we give up with an exception
+            raise RuntimeError(
+                "Multiple composite types within unions are not supported. Received: "
+                f"{', '.join(str(t.field_definition.raw) for t in non_simple_types)}"
+            )
+
+        # special case: simple + one non-simple
+        if len(non_simple_types) == 1 and simple_types:
+            conditional = "if"
+            for simple_type in simple_types:
+                _handle_transfer_instance(simple_type, conditional_=conditional)
+                conditional = "elif"
+
+            with self._start_block("else:"):
+                self._create_transfer_type_data_body(
+                    transfer_type=non_simple_types[0],
+                    nested_as_dict=False,
+                    source_value_name=source_value_name,
+                    assignment_target=assignment_target,
+                )
+            return
+
+        for inner_type in simple_types:
+            if inner_type.nested_field_info:
+                _handle_transfer_instance(inner_type, conditional_="if")
+                return
+
         self._add_stmt(f"{assignment_target} = {source_value_name}")

tests/unit/test_dto/test_factory/test_backends/test_backends.py

@@ -1,17 +1,19 @@
 from __future__ import annotations
 
 import inspect
+import re
 from dataclasses import dataclass, field
 from types import ModuleType
 from typing import TYPE_CHECKING, Callable, Optional
 from unittest.mock import MagicMock
 
+import msgspec
 import pytest
 from msgspec import Meta, Struct, to_builtins
 
 from litestar import Litestar, Request, get, post
 from litestar._openapi.schema_generation import SchemaCreator
-from litestar.dto import DataclassDTO, DTOConfig, DTOField
+from litestar.dto import DataclassDTO, DTOConfig, DTOField, MsgspecDTO
 from litestar.dto._backend import DTOBackend, _create_struct_field_meta_for_field_definition
 from litestar.dto._codegen_backend import DTOCodegenBackend
 from litestar.dto._types import CollectionType, SimpleType, TransferDTOFieldDefinition
@@ -549,3 +551,96 @@ def test_create_struct_field_meta_for_field_definition(constraint_kwargs: Any) -
         title="test",
         **constraint_kwargs,
     )
+
+
+@pytest.mark.parametrize(
+    "simple_type, value",
+    [
+        ("None", [{"value": "hello"}]),
+        ("None", None),
+        ("None,int", None),
+        ("None,int", 1),
+        ("None,int", [{"value": "hello"}]),
+        ("int", [{"value": "hello"}]),
+        ("int", 1),
+        ("bool", [{"value": "hello"}]),
+        ("bool", True),
+        ("bool,str,int", True),
+        ("bool,str,int", 1),
+        ("bool,str,int", "hello"),
+        ("bool,str,int", [{"value": "hello"}]),
+        ("bool,Inner", {"value": "hello"}),
+        ("bool,Inner", [{"value": "hello"}]),
+        ("bool,Inner", True),
+    ],
+)
+def test_transfer_nested_simple_type_union(
+    asgi_connection: Request[Any, Any, Any],
+    create_module: Callable[[str], ModuleType],
+    simple_type: str,
+    value: Any,
+) -> None:
+    # https://github.com/litestar-org/litestar/issues/4273
+
+    module = create_module(f"""
+from typing import Union
+import msgspec
+
+class Inner(msgspec.Struct):
+    value: str
+
+class Outer(msgspec.Struct):
+    some_field: Union[{simple_type}, list[Inner]]
+""")
+
+    backend = DTOCodegenBackend(
+        handler_id="test",
+        dto_factory=MsgspecDTO[module.Outer],  # type: ignore[name-defined]
+        field_definition=TransferDTOFieldDefinition.from_annotation(module.Outer),
+        model_type=module.Outer,
+        wrapper_attribute_name=None,
+        is_data_field=True,
+    )
+
+    data = backend.populate_data_from_builtins({"some_field": value}, asgi_connection)
+    assert isinstance(data, module.Outer)
+    if isinstance(value, list):
+        assert data.some_field == msgspec.convert(value, type=list[module.Inner])  # type: ignore[name-defined]
+    elif isinstance(value, dict):
+        assert data.some_field == msgspec.convert(value, type=module.Inner)
+    else:
+        assert data.some_field == value
+
+
+def test_nested_union_with_multiple_composite_types_raises(
+    asgi_connection: Request[Any, Any, Any],
+    create_module: Callable[[str], ModuleType],
+) -> None:
+    module = create_module("""
+from typing import Union
+import dataclasses
+
+@dataclasses.dataclass
+class Inner:
+    value: str
+
+
+@dataclasses.dataclass
+class Outer:
+    some_field: Union[list[str], dict[str, str], Inner]
+""")
+
+    with pytest.raises(
+        RuntimeError,
+        match=re.escape(
+            "Multiple composite types within unions are not supported. Received: list[str], dict[str, str]"
+        ),
+    ):
+        DTOCodegenBackend(
+            handler_id="test",
+            dto_factory=DataclassDTO[module.Outer],  # type: ignore[name-defined]
+            field_definition=TransferDTOFieldDefinition.from_annotation(module.Outer),
+            model_type=module.Outer,
+            wrapper_attribute_name=None,
+            is_data_field=True,
+        )