feat: support yield message callback

Author: Lancetnik

.github/workflows/test.yml

@@ -41,7 +41,6 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.8", "3.9", "3.10", "3.11"]
-        pydantic-version: ["pydantic-v1", "pydantic-v2"]
       fail-fast: false
 
     steps:
@@ -58,12 +57,6 @@ jobs:
       - name: Install Dependencies
         if: steps.cache.outputs.cache-hit != 'true'
         run: pip install -e .[test]
-      - name: Install Pydantic v1
-        if: matrix.pydantic-version == 'pydantic-v1'
-        run: pip install "pydantic>=1.10,<2"
-      - name: Install Pydantic v2
-        if: matrix.pydantic-version == 'pydantic-v2'
-        run: pip install --pre "pydantic>=2,<3"
       - run: mkdir coverage
       - name: Test
         run: bash scripts/test.sh

README.md

@@ -71,7 +71,9 @@ taskiq_broker = BrokerWrapper(broker)
 # create periodic task
 taskiq_broker.task(
     message="Hi!",
-    subject="test-subject"
+    # If you are using RabbitBroker, then you need to replace subject with queue.
+    # If you are using KafkaBroker, then you need to replace subject with topic.
+    subject="test-subject",
     schedule=[{
         "cron": "* * * * *",
     }],
@@ -122,3 +124,17 @@ taskiq_broker.task(
     ...,
 )
 ```
+
+Also, you can send a multiple message by one task call just using generator message callback with `yield`
+
+```python
+async def collect_information_to_send():
+    """Sends 10 messages per task call."""
+    for i in range(10):
+        yield i
+
+taskiq_broker.task(
+    message=collect_information_to_send,
+    ...,
+)
+```

taskiq_faststream/__about__.py

@@ -1,2 +1,2 @@
 """FastStream - taskiq integration to schedule FastStream tasks."""
-__version__ = "0.1.3"
+__version__ = "0.1.4"

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"