feat: support multiple push

Author: Lancetnik

taskiq_faststream/broker.py

@@ -133,5 +133,7 @@ async def _broker_publish(
 ) -> None:
     labels = message.labels
     labels.pop("schedule", None)
-    msg = await resolve_msg(labels.pop("message", message.message))
-    await broker.publish(msg, **labels)
+    async for msg in resolve_msg(
+        msg=labels.pop("message", message.message),
+    ):
+        await broker.publish(msg, **labels)

taskiq_faststream/utils.py

@@ -1,5 +1,6 @@
 import typing
 
+from fast_depends.utils import is_async_gen_callable, is_gen_callable
 from faststream.types import SendableMessage
 from faststream.utils.functions import to_async
 
@@ -10,8 +11,10 @@ async def resolve_msg(
         SendableMessage,
         typing.Callable[[], SendableMessage],
         typing.Callable[[], typing.Awaitable[SendableMessage]],
+        typing.Callable[[], typing.Iterator[SendableMessage]],
+        typing.Callable[[], typing.AsyncIterator[SendableMessage]],
     ],
-) -> SendableMessage:
+) -> typing.AsyncIterator[SendableMessage]:
     """Resolve message generation callback.
 
     Args:
@@ -21,9 +24,26 @@ async def resolve_msg(
         The message to send
     """
     if callable(msg):
-        get_msg = typing.cast(
-            typing.Callable[[], typing.Awaitable[SendableMessage]],
-            to_async(msg),
-        )
-        msg = await get_msg()
-    return msg
+        if is_async_gen_callable(msg):
+            async for i in typing.cast(
+                typing.Callable[[], typing.AsyncIterator[SendableMessage]],
+                msg,
+            )():
+                yield i
+
+        elif is_gen_callable(msg):
+            for i in typing.cast(
+                typing.Callable[[], typing.Iterator[SendableMessage]],
+                msg,
+            )():
+                yield i
+
+        else:
+            get_msg = typing.cast(
+                typing.Callable[[], typing.Awaitable[SendableMessage]],
+                to_async(msg),
+            )
+            yield await get_msg()
+
+    else:
+        yield msg

tests/test_resolve_message.py

@@ -1,24 +1,29 @@
+from typing import AsyncIterator, Iterator
+
 import pytest
 
 from taskiq_faststream.utils import resolve_msg
 
 
 @pytest.mark.anyio
 async def test_regular() -> None:
-    assert await resolve_msg("msg") == "msg"
+    async for m in resolve_msg("msg"):
+        assert m == "msg"
 
 
 @pytest.mark.anyio
 async def test_sync_callable() -> None:
-    assert await resolve_msg(lambda: "msg") == "msg"
+    async for m in resolve_msg(lambda: "msg"):
+        assert m == "msg"
 
 
 @pytest.mark.anyio
 async def test_async_callable() -> None:
     async def gen_msg() -> str:
         return "msg"
 
-    assert await resolve_msg(gen_msg) == "msg"
+    async for m in resolve_msg(gen_msg):
+        assert m == "msg"
 
 
 @pytest.mark.anyio
@@ -30,7 +35,8 @@ def __init__(self) -> None:
         def __call__(self) -> str:
             return "msg"
 
-    assert await resolve_msg(C()) == "msg"
+    async for m in resolve_msg(C()):
+        assert m == "msg"
 
 
 @pytest.mark.anyio
@@ -42,4 +48,23 @@ def __init__(self) -> None:
         async def __call__(self) -> str:
             return "msg"
 
-    assert await resolve_msg(C()) == "msg"
+    async for m in resolve_msg(C()):
+        assert m == "msg"
+
+
+@pytest.mark.anyio
+async def test_async_generator() -> None:
+    async def get_msg() -> AsyncIterator[str]:
+        yield "msg"
+
+    async for m in resolve_msg(get_msg):
+        assert m == "msg"
+
+
+@pytest.mark.anyio
+async def test_sync_generator() -> None:
+    def get_msg() -> Iterator[str]:
+        yield "msg"
+
+    async for m in resolve_msg(get_msg):
+        assert m == "msg"