fix: add custom serializer and scheduler

Author: Lancetnik

README.md

@@ -61,9 +61,8 @@ async def handler(msg: str):
     print(msg)
 
 # taskiq-faststream scheduling
-from taskiq import TaskiqScheduler
 from taskiq.schedule_sources import LabelScheduleSource
-from taskiq_faststream import BrokerWrapper
+from taskiq_faststream import BrokerWrapper, StreamScheduler
 
 # wrap FastStream object
 taskiq_broker = BrokerWrapper(broker)
@@ -80,7 +79,7 @@ taskiq_broker.task(
 )
 
 # create scheduler object
-scheduler = TaskiqScheduler(
+scheduler = StreamScheduler(
     broker=taskiq_broker,
     sources=[LabelScheduleSource(taskiq_broker)],
 )

taskiq_faststream/__init__.py

@@ -1,6 +1,8 @@
 from taskiq_faststream.broker import AppWrapper, BrokerWrapper
+from taskiq_faststream.scheduler import StreamScheduler
 
 __all__ = (
     "BrokerWrapper",
+    "StreamScheduler",
     "AppWrapper",
 )

taskiq_faststream/broker.py

@@ -8,8 +8,9 @@
 from faststream.types import SendableMessage
 from taskiq import AsyncBroker, BrokerMessage
 from taskiq.acks import AckableMessage
+from taskiq.decor import AsyncTaskiqDecoratedTask
 
-from taskiq_faststream.task import PatchedTaskiqDecoratedTask
+from taskiq_faststream.serializer import PatchedSerializer
 from taskiq_faststream.types import ScheduledTask
 from taskiq_faststream.utils import resolve_msg
 
@@ -32,8 +33,8 @@ class BrokerWrapper(AsyncBroker):
 
     def __init__(self, broker: BrokerAsyncUsecase[typing.Any, typing.Any]) -> None:
         super().__init__()
+        self.serializer = PatchedSerializer()
         self.broker = broker
-        self.decorator_class = PatchedTaskiqDecoratedTask
 
     async def startup(self) -> None:
         """Startup wrapped FastStream broker."""
@@ -77,7 +78,7 @@ def task(  # type: ignore[override]
         *,
         schedule: typing.List[ScheduledTask],
         **kwargs: PublishParameters,
-    ) -> "PatchedTaskiqDecoratedTask[[], None]":
+    ) -> "AsyncTaskiqDecoratedTask[[], None]":
         """Register FastStream scheduled task.
 
         Args:
@@ -108,8 +109,8 @@ class AppWrapper(BrokerWrapper):
 
     def __init__(self, app: FastStream) -> None:
         super(BrokerWrapper, self).__init__()
+        self.serializer = PatchedSerializer()
         self.app = app
-        self.decorator_class = PatchedTaskiqDecoratedTask
 
     async def startup(self) -> None:
         """Startup wrapped FastStream."""

taskiq_faststream/scheduler.py

@@ -0,0 +1,28 @@
+from typing import TYPE_CHECKING
+
+from taskiq.scheduler.scheduled_task import ScheduledTask
+from taskiq.scheduler.scheduler import TaskiqScheduler as Scheduler
+from taskiq.utils import maybe_awaitable
+
+from taskiq_faststream.kicker import LabelRespectKicker
+
+if TYPE_CHECKING:  # pragma: no cover
+    from taskiq.abc.schedule_source import ScheduleSource
+
+
+class StreamScheduler(Scheduler):
+    """Patched scheduler with custom kicker."""
+
+    async def on_ready(self, source: "ScheduleSource", task: ScheduledTask) -> None:
+        """
+        This method is called when task is ready to be enqueued.
+
+        It's triggered on proper time depending on `task.cron` or `task.time` attribute.
+        :param task: task to send
+        """
+        await maybe_awaitable(source.pre_send(task))
+        await LabelRespectKicker(task.task_name, self.broker, task.labels).kiq(
+            *task.args,
+            **task.kwargs,
+        )
+        await maybe_awaitable(source.post_send(task))

taskiq_faststream/serializer.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from taskiq.serializers.json_serializer import JSONSerializer
+
+
+class PatchedSerializer(JSONSerializer):
+    """Patched serializer removes labels."""
+
+    def dumpb(self, value: Any) -> bytes:
+        """
+        Dumps taskiq message to some broker message format.
+
+        :param message: message to send.
+        :return: Dumped message.
+        """
+        del value["labels"]
+        return super().dumpb(value)