Implement customizable serializer (#214)

Author: asvetlov

poetry.lock

@@ -61,6 +61,7 @@ freezegun = "^1.2.2"
 pytest-mock = "^3.11.1"
 tzlocal = "^5.0.1"
 types-tzlocal = "^5.0.1.1"
+types-pytz = "^2023.3.1.1"
 
 [tool.poetry.extras]
 zmq = ["pyzmq"]

pyproject.toml

@@ -25,12 +25,14 @@
 from typing_extensions import ParamSpec, Self, TypeAlias
 
 from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.abc.serializer import TaskiqSerializer
 from taskiq.acks import AckableMessage
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.events import TaskiqEvents
-from taskiq.formatters.json_formatter import JSONFormatter
+from taskiq.formatters.proxy_formatter import ProxyFormatter
 from taskiq.message import BrokerMessage
 from taskiq.result_backends.dummy import DummyResultBackend
+from taskiq.serializers.json_serializer import JSONSerializer
 from taskiq.state import TaskiqState
 from taskiq.utils import maybe_awaitable, remove_suffix
 from taskiq.warnings import TaskiqDeprecationWarning
@@ -97,7 +99,8 @@ def __init__(
         self.middlewares: "List[TaskiqMiddleware]" = []
         self.result_backend = result_backend
         self.decorator_class = AsyncTaskiqDecoratedTask
-        self.formatter: "TaskiqFormatter" = JSONFormatter()
+        self.serializer: TaskiqSerializer = JSONSerializer()
+        self.formatter: "TaskiqFormatter" = ProxyFormatter(self)
         self.id_generator = task_id_generator
         self.local_task_registry: Dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {}
         # Every event has a list of handlers.
@@ -479,6 +482,19 @@ def with_event_handlers(
         self.event_handlers[event].extend(handlers)
         return self
 
+    def with_serializer(
+        self,
+        serializer: TaskiqSerializer,
+    ) -> "Self":  # pragma: no cover
+        """
+        Set a new serializer and return an updated broker.
+
+        :param serializer: new serializer.
+        :return: self
+        """
+        self.serializer = serializer
+        return self
+
     def _register_task(
         self,
         task_name: str,

taskiq/abc/broker.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class TaskiqSerializer(ABC):
+    """Custom serializer for brokers."""
+
+    @abstractmethod
+    def dumpb(self, value: Any) -> bytes:
+        """
+        Dump value to bytes for sending through the wire.
+
+        :param value: value to encode.
+        :return: encoded value.
+        """
+
+    @abstractmethod
+    def loadb(self, value: bytes) -> Any:
+        """
+        Parse byte-encoded value received from the wire.
+
+        :param message: value to parse.
+        :return: decoded value.
+        """

taskiq/abc/serializer.py

@@ -16,6 +16,15 @@
     def parse_obj_as(annot: T, obj: Any) -> T:
         return pydantic.TypeAdapter(annot).validate_python(obj)
 
+    def model_validate(
+        model_class: Type[Model],
+        message: Dict[str, Any],
+    ) -> Model:
+        return model_class.model_validate(message)
+
+    def model_dump(instance: Model) -> Dict[str, Any]:
+        return instance.model_dump()
+
     def model_validate_json(
         model_class: Type[Model],
         message: Union[str, bytes, bytearray],
@@ -37,6 +46,15 @@ def model_copy(
 else:
     parse_obj_as = pydantic.parse_obj_as  # type: ignore
 
+    def model_validate(
+        model_class: Type[Model],
+        message: Dict[str, Any],
+    ) -> Model:
+        return model_class.parse_obj(message)
+
+    def model_dump(instance: Model) -> Dict[str, Any]:
+        return instance.dict()
+
     def model_validate_json(
         model_class: Type[Model],
         message: Union[str, bytes, bytearray],

taskiq/compat.py

@@ -4,7 +4,7 @@
 
 
 class JSONFormatter(TaskiqFormatter):
-    """Default taskiq formatter."""
+    """JSON taskiq formatter."""
 
     def dumps(self, message: TaskiqMessage) -> BrokerMessage:
         """

taskiq/formatters/json_formatter.py

@@ -0,0 +1,38 @@
+from typing import TYPE_CHECKING
+
+from taskiq.abc.formatter import TaskiqFormatter
+from taskiq.compat import model_dump, model_validate
+from taskiq.message import BrokerMessage, TaskiqMessage
+
+if TYPE_CHECKING:
+    from taskiq.abc.broker import AsyncBroker
+
+
+class ProxyFormatter(TaskiqFormatter):
+    """Default taskiq formatter."""
+
+    def __init__(self, broker: "AsyncBroker") -> None:
+        self.broker = broker
+
+    def dumps(self, message: TaskiqMessage) -> BrokerMessage:
+        """
+        Dumps taskiq message to some broker message format.
+
+        :param message: message to send.
+        :return: Dumped message.
+        """
+        return BrokerMessage(
+            task_id=message.task_id,
+            task_name=message.task_name,
+            message=self.broker.serializer.dumpb(model_dump(message)),
+            labels=message.labels,
+        )
+
+    def loads(self, message: bytes) -> TaskiqMessage:
+        """
+        Loads json from message.
+
+        :param message: broker's message.
+        :return: parsed taskiq message.
+        """
+        return model_validate(TaskiqMessage, self.broker.serializer.loadb(message))

taskiq/formatters/proxy_formatter.py

@@ -14,7 +14,7 @@ class TaskiqMessage(BaseModel):
 
     task_id: str
     task_name: str
-    labels: Dict[str, str]
+    labels: Dict[str, Any]
     args: List[Any]
     kwargs: Dict[str, Any]
 
@@ -25,4 +25,4 @@ class BrokerMessage(BaseModel):
     task_id: str
     task_name: str
     message: bytes
-    labels: Dict[str, str]
+    labels: Dict[str, Any]

taskiq/message.py

@@ -0,0 +1 @@
+"""Taskiq serializers."""

taskiq/serializers/__init__.py

@@ -0,0 +1,26 @@
+from json import dumps, loads
+from typing import Any
+
+from taskiq.abc.serializer import TaskiqSerializer
+
+
+class JSONSerializer(TaskiqSerializer):
+    """Default taskiq serizalizer."""
+
+    def dumpb(self, value: Any) -> bytes:
+        """
+        Dumps taskiq message to some broker message format.
+
+        :param message: message to send.
+        :return: Dumped message.
+        """
+        return dumps(value).encode()
+
+    def loadb(self, value: bytes) -> Any:
+        """
+        Parse byte-encoded value received from the wire.
+
+        :param message: value to parse.
+        :return: decoded value.
+        """
+        return loads(value.decode())