Updated broker's interface.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

poetry.lock

@@ -19,7 +19,7 @@ keywords = ["taskiq", "tasks", "distributed", "async", "redis", "result_backend"
 
 [tool.poetry.dependencies]
 python = "^3.7"
-taskiq = "^0"
+taskiq = "^0.2.0"
 redis = "^4.2.0"
 
 [tool.poetry.dev-dependencies]

pyproject.toml

@@ -1,7 +1,6 @@
-import asyncio
 import pickle
 from logging import getLogger
-from typing import Any, Callable, Coroutine, Optional, TypeVar
+from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
 
 from redis.asyncio import ConnectionPool, Redis
 from taskiq.abc.broker import AsyncBroker
@@ -70,19 +69,15 @@ async def kick(self, message: BrokerMessage) -> None:
                 pickle.dumps(message),
             )
 
-    async def listen(
-        self,
-        callback: Callable[[BrokerMessage], Coroutine[Any, Any, None]],
-    ) -> None:
+    async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
         """
-        Listen redis list for new messages.
+        Listen redis queue for new messages.
 
-        This function listens to list calls callback on
-        new messages.
+        This function listens to the queue
+        and yields new messages if they have BrokerMessage type.
 
-        :param callback: function to call on new message.
+        :yields: broker messages.
         """
-        loop = asyncio.get_event_loop()
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
             redis_pubsub_channel = redis_conn.pubsub()
             await redis_pubsub_channel.subscribe(self.redis_pubsub_channel)
@@ -93,7 +88,7 @@ async def listen(
                             message["data"],
                         )
                         if isinstance(redis_message, BrokerMessage):
-                            loop.create_task(callback(redis_message))
+                            yield redis_message
                     except (
                         TypeError,
                         AttributeError,