feat (#1447): add StreamMessage.batch_headers attr to provide access to whole batch messages headers

Author: Lancetnik

faststream/broker/message.py

@@ -6,6 +6,7 @@
     TYPE_CHECKING,
     Any,
     Generic,
+    List,
     Optional,
     Sequence,
     Tuple,
@@ -38,6 +39,7 @@ class StreamMessage(Generic[MsgType]):
 
     body: Union[bytes, Any]
     headers: "AnyDict" = field(default_factory=dict)
+    batch_headers: List["AnyDict"] = field(default_factory=list)
     path: "AnyDict" = field(default_factory=dict)
 
     content_type: Optional[str] = None

faststream/broker/publisher/proto.py

@@ -57,8 +57,7 @@ class PublisherProto(
     _producer: Optional["ProducerProto"]
 
     @abstractmethod
-    def add_middleware(self, middleware: "BrokerMiddleware[MsgType]") -> None:
-        ...
+    def add_middleware(self, middleware: "BrokerMiddleware[MsgType]") -> None: ...
 
     @staticmethod
     @abstractmethod

faststream/broker/subscriber/proto.py

@@ -36,8 +36,7 @@ class SubscriberProto(
     _producer: Optional["ProducerProto"]
 
     @abstractmethod
-    def add_middleware(self, middleware: "BrokerMiddleware[MsgType]") -> None:
-        ...
+    def add_middleware(self, middleware: "BrokerMiddleware[MsgType]") -> None: ...
 
     @staticmethod
     @abstractmethod

faststream/cli/docs/app.py

@@ -45,8 +45,7 @@ def serve(
         ),
     ),
     is_factory: bool = typer.Option(
-        False,
-        "--factory", help="Treat APP as an application factory"
+        False, "--factory", help="Treat APP as an application factory"
     ),
 ) -> None:
     """Serve project AsyncAPI schema."""
@@ -109,8 +108,7 @@ def gen(
         ),
     ),
     is_factory: bool = typer.Option(
-        False,
-        "--factory", help="Treat APP as an application factory"
+        False, "--factory", help="Treat APP as an application factory"
     ),
 ) -> None:
     """Generate project AsyncAPI schema."""

faststream/cli/main.py

@@ -210,8 +210,7 @@ def publish(
     message: str = typer.Argument(..., help="Message to be published"),
     rpc: bool = typer.Option(False, help="Enable RPC mode and system output"),
     is_factory: bool = typer.Option(
-        False,
-        "--factory", help="Treat APP as an application factory"
+        False, "--factory", help="Treat APP as an application factory"
     ),
 ) -> None:
     """Publish a message using the specified broker in a FastStream application.

faststream/kafka/parser.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 from faststream.broker.message import decode_message, gen_cor_id
 from faststream.kafka.message import FAKE_CONSUMER, KafkaMessage
@@ -39,13 +39,24 @@ async def parse_message_batch(
         message: Tuple["ConsumerRecord", ...],
     ) -> "StreamMessage[Tuple[ConsumerRecord, ...]]":
         """Parses a batch of messages from a Kafka consumer."""
+        body: List[Any] = []
+        batch_headers: List[Dict[str, str]] = []
+
         first = message[0]
         last = message[-1]
-        headers = {i: j.decode() for i, j in first.headers}
+
+        for m in message:
+            body.append(m.value)
+            batch_headers.append({i: j.decode() for i, j in m.headers})
+
+        headers = next(iter(batch_headers), {})
+
         handler: Optional["LogicSubscriber[Any]"] = context.get_local("handler_")
+
         return KafkaMessage(
-            body=[m.value for m in message],
+            body=body,
             headers=headers,
+            batch_headers=batch_headers,
             reply_to=headers.get("reply_to", ""),
             content_type=headers.get("content-type"),
             message_id=f"{first.offset}-{last.offset}-{first.timestamp}",

faststream/nats/parser.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional
 
 from faststream.broker.message import StreamMessage, decode_message, gen_cor_id
 from faststream.nats.message import NatsBatchMessage, NatsMessage
@@ -102,19 +102,27 @@ async def parse_batch(
         self,
         message: List["Msg"],
     ) -> "StreamMessage[List[Msg]]":
-        if first_msg := next(iter(message), None):
-            path = self.get_path(first_msg.subject)
-            headers = first_msg.headers
+        body: List[bytes] = []
+        batch_headers: List[Dict[str, str]] = []
+
+        if message:
+            path = self.get_path(message[0].subject)
+
+            for m in message:
+                batch_headers.append(m.headers or {})
+                body.append(m.data)
 
         else:
             path = None
-            headers = None
+
+        headers = next(iter(batch_headers), {})
 
         return NatsBatchMessage(
             raw_message=message,
-            body=[m.data for m in message],
+            body=body,
             path=path or {},
-            headers=headers or {},
+            headers=headers,
+            batch_headers=batch_headers,
         )
 
     async def decode_batch(

faststream/rabbit/broker/broker.py

@@ -118,7 +118,7 @@ def __init__(
             Doc(
                 "raise an :class:`aio_pika.exceptions.DeliveryError`"
                 "when mandatory message will be returned"
-            )
+            ),
         ] = False,
         # broker args
         max_consumers: Annotated[
@@ -345,7 +345,7 @@ async def connect(  # type: ignore[override]
             Doc(
                 "raise an :class:`aio_pika.exceptions.DeliveryError`"
                 "when mandatory message will be returned"
-            )
+            ),
         ] = Parameter.empty,
     ) -> "RobustConnection":
         """Connect broker object to RabbitMQ.

faststream/rabbit/fastapi/router.py

@@ -114,7 +114,7 @@ def __init__(
             Doc(
                 "raise an :class:`aio_pika.exceptions.DeliveryError`"
                 "when mandatory message will be returned"
-            )
+            ),
         ] = False,
         # broker args
         max_consumers: Annotated[

faststream/redis/parser.py

@@ -1,7 +1,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    Dict,
+    List,
     Mapping,
     Optional,
     Sequence,
@@ -136,22 +136,27 @@ async def parse_message(
         self,
         message: Mapping[str, Any],
     ) -> "StreamMessage[Mapping[str, Any]]":
-        data, headers = self._parse_data(message)
+        data, headers, batch_headers = self._parse_data(message)
+
         id_ = gen_cor_id()
+
         return self.msg_class(
             raw_message=message,
             body=data,
             path=self.get_path(message),
             headers=headers,
+            batch_headers=batch_headers,
             reply_to=headers.get("reply_to", ""),
             content_type=headers.get("content-type"),
             message_id=headers.get("message_id", id_),
             correlation_id=headers.get("correlation_id", id_),
         )
 
     @staticmethod
-    def _parse_data(message: Mapping[str, Any]) -> Tuple[bytes, "AnyDict"]:
-        return RawMessage.parse(message["data"])
+    def _parse_data(
+        message: Mapping[str, Any],
+    ) -> Tuple[bytes, "AnyDict", List["AnyDict"]]:
+        return (*RawMessage.parse(message["data"]), [])
 
     def get_path(self, message: Mapping[str, Any]) -> "AnyDict":
         if (
@@ -183,42 +188,68 @@ class RedisBatchListParser(SimpleParser):
     msg_class = RedisBatchListMessage
 
     @staticmethod
-    def _parse_data(message: Mapping[str, Any]) -> Tuple[bytes, "AnyDict"]:
-        data = [_decode_batch_body_item(x) for x in message["data"]]
+    def _parse_data(
+        message: Mapping[str, Any],
+    ) -> Tuple[bytes, "AnyDict", List["AnyDict"]]:
+        body: List[Any] = []
+        batch_headers: List["AnyDict"] = []
+
+        for x in message["data"]:
+            msg_data, msg_headers = _decode_batch_body_item(x)
+            body.append(msg_data)
+            batch_headers.append(msg_headers)
+
+        first_msg_headers = next(iter(batch_headers), {})
+
         return (
-            dump_json(i[0] for i in data),
+            dump_json(body),
             {
-                **data[0][1],
-                "content-type": ContentTypes.json,
+                **first_msg_headers,
+                "content-type": ContentTypes.json.value,
             },
+            batch_headers,
         )
 
 
 class RedisStreamParser(SimpleParser):
     msg_class = RedisStreamMessage
 
     @classmethod
-    def _parse_data(cls, message: Mapping[str, Any]) -> Tuple[bytes, "AnyDict"]:
+    def _parse_data(
+        cls, message: Mapping[str, Any]
+    ) -> Tuple[bytes, "AnyDict", List["AnyDict"]]:
         data = message["data"]
-        return RawMessage.parse(data.get(bDATA_KEY) or dump_json(data))
+        return (*RawMessage.parse(data.get(bDATA_KEY) or dump_json(data)), [])
 
 
 class RedisBatchStreamParser(SimpleParser):
     msg_class = RedisBatchStreamMessage
 
     @staticmethod
-    def _parse_data(message: Mapping[str, Any]) -> Tuple[bytes, "AnyDict"]:
-        data = [_decode_batch_body_item(x.get(bDATA_KEY, x)) for x in message["data"]]
+    def _parse_data(
+        message: Mapping[str, Any],
+    ) -> Tuple[bytes, "AnyDict", List["AnyDict"]]:
+        body: List[Any] = []
+        batch_headers: List["AnyDict"] = []
+
+        for x in message["data"]:
+            msg_data, msg_headers = _decode_batch_body_item(x.get(bDATA_KEY, x))
+            body.append(msg_data)
+            batch_headers.append(msg_headers)
+
+        first_msg_headers = next(iter(batch_headers), {})
+
         return (
-            dump_json(i[0] for i in data),
+            dump_json(body),
             {
-                **data[0][1],
-                "content-type": ContentTypes.json,
+                **first_msg_headers,
+                "content-type": ContentTypes.json.value,
             },
+            batch_headers,
         )
 
 
-def _decode_batch_body_item(msg_content: bytes) -> Tuple[Any, Dict[str, str]]:
+def _decode_batch_body_item(msg_content: bytes) -> Tuple[Any, "AnyDict"]:
     msg_body, headers = RawMessage.parse(msg_content)
     try:
         return json_loads(msg_body), headers

faststream/types.py

@@ -63,22 +63,16 @@ class StandardDataclass(Protocol):
     """Protocol to check type is dataclass."""
 
     __dataclass_fields__: ClassVar[Dict[str, Any]]
-    __dataclass_params__: ClassVar[Any]
-    __post_init__: ClassVar[Callable[..., None]]
-
-    def __init__(self, *args: object, **kwargs: object) -> None:
-        """Interface method."""
-        ...
 
 
 BaseSendableMessage: TypeAlias = Union[
     JsonDecodable,
     Decimal,
     datetime,
-    None,
     StandardDataclass,
     SendableTable,
     SendableArray,
+    None,
 ]
 
 try:

tests/brokers/base/middlewares.py

@@ -271,7 +271,11 @@ async def handler(m):
         mock.end.assert_called_once()
 
     async def test_add_global_middleware(
-        self, event: asyncio.Event, queue: str, mock: Mock, raw_broker,
+        self,
+        event: asyncio.Event,
+        queue: str,
+        mock: Mock,
+        raw_broker,
     ):
         class mid(BaseMiddleware):  # noqa: N801
             async def on_receive(self):

tests/brokers/base/publish.py

@@ -1,4 +1,5 @@
 import asyncio
+from dataclasses import asdict, dataclass
 from datetime import datetime
 from typing import Any, ClassVar, Dict, List, Tuple
 from unittest.mock import Mock
@@ -7,7 +8,7 @@
 import pytest
 from pydantic import BaseModel
 
-from faststream._compat import model_to_json
+from faststream._compat import dump_json, model_to_json
 from faststream.annotations import Logger
 from faststream.broker.core.usecase import BrokerUsecase
 
@@ -16,6 +17,11 @@ class SimpleModel(BaseModel):
     r: str
 
 
+@dataclass
+class SimpleDataclass:
+    r: str
+
+
 now = datetime.now()
 
 
@@ -55,6 +61,12 @@ def pub_broker(self, full_broker):
                 1.0,
                 id="float->float",
             ),
+            pytest.param(
+                1,
+                float,
+                1.0,
+                id="int->float",
+            ),
             pytest.param(
                 False,
                 bool,
@@ -103,6 +115,30 @@ def pub_broker(self, full_broker):
                 SimpleModel(r="hello!"),
                 id="dict->model",
             ),
+            pytest.param(
+                dump_json(asdict(SimpleDataclass(r="hello!"))),
+                SimpleDataclass,
+                SimpleDataclass(r="hello!"),
+                id="bytes->dataclass",
+            ),
+            pytest.param(
+                SimpleDataclass(r="hello!"),
+                SimpleDataclass,
+                SimpleDataclass(r="hello!"),
+                id="dataclass->dataclass",
+            ),
+            pytest.param(
+                SimpleDataclass(r="hello!"),
+                dict,
+                {"r": "hello!"},
+                id="dataclass->dict",
+            ),
+            pytest.param(
+                {"r": "hello!"},
+                SimpleDataclass,
+                SimpleDataclass(r="hello!"),
+                id="dict->dataclass",
+            ),
         ),
     )
     async def test_serialize(