Move UnionModel logic to decode_message

Author: WiZeYAR

src/asyncapi_python/amqp/endpoint/receiver.py

@@ -22,7 +22,6 @@
 from .base import AbstractEndpoint, EndpointParams
 from ..error import Rejection, BadRequestRejection
 from ..operation import Operation
-from ..utils import UnionModel
 from aio_pika.abc import AbstractIncomingMessage, AbstractRobustQueue
 
 
@@ -81,9 +80,6 @@ async def _consumer(self, message: AbstractIncomingMessage):
     def _decode_payload(self, message: AbstractIncomingMessage) -> I:
         try:
             payload: I = self._params.decode(message.body, self._op.message_type)
-            # This trick is used to support union types in receivers
-            if isinstance(payload, UnionModel):
-                payload = cast(I, payload.root)
         except ValidationError as e:
             raise BadRequestRejection(e)
         return payload

src/asyncapi_python/amqp/utils.py

@@ -15,21 +15,25 @@
 
 from functools import cache
 from pydantic import BaseModel, RootModel
-from typing import TypeVar, Union
+from typing import TypeVar, Union, cast
 
 T = TypeVar("T", bound=BaseModel)
 U = TypeVar("U")
 
 
-class UnionModel(RootModel[U]): ...
+class UnionModel(RootModel[U]):
+    """A trick to allow unions as constructor types"""
 
 
 def encode_message(message: T) -> bytes:
     return message.model_dump_json().encode()
 
 
 def decode_message(message: bytes, schema: type[T]) -> T:
-    return schema.model_validate_json(message)
+    payload = schema.model_validate_json(message)
+    if isinstance(payload, UnionModel):
+        payload = cast(T, payload.root)
+    return payload
 
 
 @cache