divided redis message format into different classes and added arguments for choosing format

Author: ilya-4real

faststream/redis/broker/broker.py

@@ -32,6 +32,7 @@
 from faststream.exceptions import NOT_CONNECTED_YET
 from faststream.redis.broker.logging import RedisLoggingBroker
 from faststream.redis.broker.registrator import RedisRegistrator
+from faststream.redis.parser import JSONMessageFormat, MessageFormat
 from faststream.redis.publisher.producer import RedisFastProducer
 from faststream.redis.security import parse_security
 from faststream.types import EMPTY
@@ -139,6 +140,10 @@ def __init__(
             Sequence["BrokerMiddleware[BaseMessage]"],
             Doc("Middlewares to apply to all broker publishers/subscribers."),
         ] = (),
+        message_format: Annotated[
+            Type["MessageFormat"],
+            Doc("What format to use when parsing messages"),
+        ] = JSONMessageFormat,
         # AsyncAPI args
         security: Annotated[
             Optional["BaseSecurity"],
@@ -202,6 +207,7 @@ def __init__(
         ] = (),
     ) -> None:
         self._producer = None
+        self.message_format = message_format
 
         if asyncapi_url is None:
             asyncapi_url = url
@@ -348,6 +354,7 @@ async def _connect(  # type: ignore[override]
             connection=client,
             parser=self._parser,
             decoder=self._decoder,
+            message_format=self.message_format,
         )
         return client
 

faststream/redis/broker/registrator.py

@@ -1,11 +1,22 @@
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Sequence, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+    cast,
+)
 
 from typing_extensions import Annotated, Doc, deprecated, override
 
 from faststream.broker.core.abc import ABCBroker
 from faststream.broker.utils import default_filter
 from faststream.exceptions import SetupError
 from faststream.redis.message import UnifyRedisDict
+from faststream.redis.parser import JSONMessageFormat, MessageFormat
 from faststream.redis.publisher.asyncapi import AsyncAPIPublisher
 from faststream.redis.subscriber.asyncapi import AsyncAPISubscriber
 from faststream.redis.subscriber.factory import SubsciberType, create_subscriber
@@ -97,6 +108,10 @@ def subscriber(  # type: ignore[override]
                 "Whether to disable **FastStream** RPC and Reply To auto responses or not."
             ),
         ] = False,
+        message_format: Annotated[
+            Type["MessageFormat"],
+            Doc("What format to use when parsing messages"),
+        ] = JSONMessageFormat,
         # AsyncAPI information
         title: Annotated[
             Optional[str],
@@ -122,6 +137,7 @@ def subscriber(  # type: ignore[override]
                     list=list,
                     stream=stream,
                     # subscriber args
+                    message_format=message_format,
                     no_ack=no_ack,
                     no_reply=no_reply,
                     retry=retry,

faststream/redis/parser.py

@@ -1,5 +1,5 @@
 import enum
-import warnings
+from abc import ABC, abstractmethod
 from struct import pack, unpack
 from typing import (
     TYPE_CHECKING,
@@ -92,18 +92,14 @@ def read_bytes(self) -> bytes:
         return self.data[self.offset :]
 
 
-class RawMessage:
+class MessageFormat(ABC):
     """A class to represent a raw Redis message."""
 
     __slots__ = (
         "data",
         "headers",
     )
 
-    IDENTITY_HEADER = (
-        b"\x89BIN\x0d\x0a\x1a\x0a"  # to avoid confusion with other formats
-    )
-
     def __init__(
         self,
         data: bytes,
@@ -120,7 +116,7 @@ def build(
         reply_to: Optional[str],
         headers: Optional["AnyDict"],
         correlation_id: str,
-    ) -> "RawMessage":
+    ) -> "MessageFormat":
         payload, content_type = encode_message(message)
 
         headers_to_send = {
@@ -141,6 +137,25 @@ def build(
             headers=headers_to_send,
         )
 
+    @classmethod
+    @abstractmethod
+    def encode(
+        cls,
+        *,
+        message: Union[Sequence["SendableMessage"], "SendableMessage"],
+        reply_to: Optional[str],
+        headers: Optional["AnyDict"],
+        correlation_id: str,
+    ) -> bytes:
+        raise NotImplementedError()
+
+    @classmethod
+    @abstractmethod
+    def parse(cls, data: bytes) -> Tuple[bytes, "AnyDict"]:
+        raise NotImplementedError()
+
+
+class JSONMessageFormat(MessageFormat):
     @classmethod
     def encode(
         cls,
@@ -149,29 +164,54 @@ def encode(
         reply_to: Optional[str],
         headers: Optional["AnyDict"],
         correlation_id: str,
-        to_json: bool = False,
     ) -> bytes:
         msg = cls.build(
             message=message,
             reply_to=reply_to,
             headers=headers,
             correlation_id=correlation_id,
         )
-        if to_json:
-            warnings.warn(
-                "JSON encoding deprecated in **FastStream 0.6.0**. "
-                "Please don't use it unless it's necessary"
-                "Format will be removed in **FastStream 0.7.0**.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-            return dump_json(
-                {
-                    "data": msg.data,
-                    "headers": msg.headers,
-                }
-            )
+        return dump_json(
+            {
+                "data": msg.data,
+                "headers": msg.headers,
+            }
+        )
+
+    @classmethod
+    def parse(cls, data: bytes) -> Tuple[bytes, "AnyDict"]:
+        headers: AnyDict
+        try:
+            parsed_data = json_loads(data)
+            data = parsed_data["data"].encode()
+            headers = parsed_data["headers"]
+        except Exception:
+            # Raw Redis message format
+            data = data
+            headers = {}
+        return data, headers
+
+
+class BinaryMessageFormatV1(MessageFormat):
+    IDENTITY_HEADER = (
+        b"\x89BIN\x0d\x0a\x1a\x0a"  # to avoid confusion with other formats
+    )
 
+    @classmethod
+    def encode(
+        cls,
+        *,
+        message: Union[Sequence["SendableMessage"], "SendableMessage"],
+        reply_to: Optional[str],
+        headers: Optional["AnyDict"],
+        correlation_id: str,
+    ) -> bytes:
+        msg = cls.build(
+            message=message,
+            reply_to=reply_to,
+            headers=headers,
+            correlation_id=correlation_id,
+        )
         writer = BinaryWriter()
         writer.write(cls.IDENTITY_HEADER)
         writer.write_int(FastStreamMessageVersion.v1.value)
@@ -185,41 +225,26 @@ def encode(
     @classmethod
     def parse(cls, data: bytes) -> Tuple[bytes, "AnyDict"]:
         headers: AnyDict
-
-        # FastStream message format
         try:
             reader = BinaryReader(data)
+            headers = {}
             magic_header = reader.read_until(len(cls.IDENTITY_HEADER))
             message_version = reader.read_int()
             if (
                 magic_header == cls.IDENTITY_HEADER
                 and message_version == FastStreamMessageVersion.v1.value
             ):
                 header_count = reader.read_int()
-                headers = {}
                 for _ in range(header_count):
                     key = reader.read_string()
                     value = reader.read_string()
                     headers[key] = value
 
                 data = reader.read_bytes()
-            else:
-                # JSON message format
-                parsed_data = json_loads(data)
-                data = parsed_data["data"].decode()
-                headers = parsed_data["headers"]
-                warnings.warn(
-                    "JSON decoding deprecated in **FastStream 0.6.0**. "
-                    "Please don't use it unless it's necessary"
-                    "Format will be removed in **FastStream 0.7.0**.",
-                    DeprecationWarning,
-                    stacklevel=2,
-                )
         except Exception:
             # Raw Redis message format
             data = data
             headers = {}
-
         return data, headers
 
 
@@ -228,9 +253,11 @@ class SimpleParser:
 
     def __init__(
         self,
+        message_format: Type["MessageFormat"] = JSONMessageFormat,
         pattern: Optional["Pattern[str]"] = None,
     ) -> None:
         self.pattern = pattern
+        self.message_format = message_format
 
     async def parse_message(
         self,
@@ -256,7 +283,7 @@ def _parse_data(
         self,
         message: Mapping[str, Any],
     ) -> Tuple[bytes, "AnyDict", List["AnyDict"]]:
-        return (*RawMessage.parse(message["data"]), [])
+        return (*self.message_format.parse(message["data"]), [])
 
     def get_path(self, message: Mapping[str, Any]) -> "AnyDict":
         if (
@@ -295,7 +322,7 @@ def _parse_data(
         batch_headers: List[AnyDict] = []
 
         for x in message["data"]:
-            msg_data, msg_headers = _decode_batch_body_item(x)
+            msg_data, msg_headers = _decode_batch_body_item(x, self.message_format)
             body.append(msg_data)
             batch_headers.append(msg_headers)
 
@@ -319,7 +346,7 @@ 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 (*JSONMessageFormat.parse(data.get(bDATA_KEY) or dump_json(data)), [])
 
 
 class RedisBatchStreamParser(SimpleParser):
@@ -333,7 +360,9 @@ def _parse_data(
         batch_headers: List[AnyDict] = []
 
         for x in message["data"]:
-            msg_data, msg_headers = _decode_batch_body_item(x.get(bDATA_KEY, x))
+            msg_data, msg_headers = _decode_batch_body_item(
+                x.get(bDATA_KEY, x), self.message_format
+            )
             body.append(msg_data)
             batch_headers.append(msg_headers)
 
@@ -349,8 +378,10 @@ def _parse_data(
         )
 
 
-def _decode_batch_body_item(msg_content: bytes) -> Tuple[Any, "AnyDict"]:
-    msg_body, headers = RawMessage.parse(msg_content)
+def _decode_batch_body_item(
+    msg_content: bytes, message_format: Type["MessageFormat"]
+) -> Tuple[Any, "AnyDict"]:
+    msg_body, headers = message_format.parse(msg_content)
     try:
         return json_loads(msg_body), headers
     except Exception:

faststream/redis/publisher/producer.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, Type
 
 import anyio
 from typing_extensions import override
@@ -7,7 +7,7 @@
 from faststream.broker.utils import resolve_custom_func
 from faststream.exceptions import WRONG_PUBLISH_ARGS, SetupError
 from faststream.redis.message import DATA_KEY
-from faststream.redis.parser import RawMessage, RedisPubSubParser
+from faststream.redis.parser import JSONMessageFormat, MessageFormat, RedisPubSubParser
 from faststream.redis.schemas import INCORRECT_SETUP_MSG
 from faststream.utils.functions import timeout_scope
 from faststream.utils.nuid import NUID
@@ -34,10 +34,12 @@ def __init__(
         connection: "Redis[bytes]",
         parser: Optional["CustomCallable"],
         decoder: Optional["CustomCallable"],
+        message_format: Type["MessageFormat"] = JSONMessageFormat,
     ) -> None:
         self._connection = connection
+        self.message_format = message_format
 
-        default = RedisPubSubParser()
+        default = RedisPubSubParser(message_format=message_format)
         self._parser = resolve_custom_func(
             parser,
             default.parse_message,
@@ -83,7 +85,7 @@ async def publish(  # type: ignore[override]
             psub = self._connection.pubsub()
             await psub.subscribe(reply_to)
 
-        msg = RawMessage.encode(
+        msg = self.message_format.encode(
             message=message,
             reply_to=reply_to,
             headers=headers,
@@ -154,7 +156,7 @@ async def request(  # type: ignore[override]
         psub = self._connection.pubsub()
         await psub.subscribe(reply_to)
 
-        msg = RawMessage.encode(
+        msg = self.message_format.encode(
             message=message,
             reply_to=reply_to,
             headers=headers,
@@ -204,7 +206,7 @@ async def publish_batch(
         pipeline: Optional["Pipeline[bytes]"] = None,
     ) -> None:
         batch = (
-            RawMessage.encode(
+            self.message_format.encode(
                 message=msg,
                 correlation_id=correlation_id,
                 reply_to=None,