Merge pull request #56 from taskiq-python/feat-get-rid-of-aiostream

feat: remove aiostream as a dependency

Author: danfimov

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2022-2025 Pavel Kirilin
+Copyright (c) 2022-2026 Pavel Kirilin
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

Makefile

@@ -12,3 +12,11 @@ help: ## Show this help
 .PHONY: clear_rabbit
 clear_rabbit:  ## Clear RabbitMQ data volume and restart container
 	@docker stop taskiq_aio_pika_rabbitmq && docker rm taskiq_aio_pika_rabbitmq &&  docker volume rm taskiq-aio-pika_rabbitmq_data && docker compose up -d
+
+.PHONY: lint
+lint:  ## Run linting
+	@uv run ruff check taskiq_aio_pika tests --fix
+
+.PHONY: run_infra
+run_infra:  ## Run infrastructure
+	@docker compose up -d rabbitmq redis

pyproject.toml

@@ -32,7 +32,7 @@ requires-python = ">=3.10,<4"
 dependencies = [
     "taskiq>=0.12.0,<1",
     "aio-pika>=9.0.0",
-    "aiostream>=0.7.1",
+    "typing-extensions>=4.14.0",
 ]
 
 [dependency-groups]
@@ -119,7 +119,7 @@ lint.ignore = [
 ]
 lint.mccabe = { max-complexity = 10 }
 exclude = [".venv/"]
-line-length = 88
+line-length = 120
 
 [tool.ruff.lint.per-file-ignores]
 "tests/*" = [

taskiq_aio_pika/broker.py

@@ -7,7 +7,6 @@
 import aiormq
 from aio_pika import DeliveryMode, ExchangeType, Message, connect_robust
 from aio_pika.abc import AbstractChannel, AbstractQueue, AbstractRobustConnection
-from aiostream import stream
 from pamqp.common import FieldTable
 from taskiq import AckableMessage, AsyncBroker, AsyncResultBackend, BrokerMessage
 from typing_extensions import Self
@@ -20,6 +19,7 @@
 )
 from taskiq_aio_pika.exchange import Exchange
 from taskiq_aio_pika.queue import Queue
+from taskiq_aio_pika.utils import merge_async_iterables
 
 _T = TypeVar("_T")
 
@@ -69,25 +69,20 @@ def __init__(
         """
         Construct a new broker.
 
-        :param url: url to rabbitmq. If None,
-            the default "amqp://guest:guest@localhost:5672" is used.
+        :param url: url to rabbitmq. If None, the default "amqp://guest:guest@localhost:5672" is used.
         :param result_backend: custom result backend.
         :param task_id_generator: custom task_id generator.
         :param qos: number of messages that worker can prefetch.
         :param loop: specific even loop.
         :param exchange: parameters of exchange that used to send messages.
-        :param task_queues: parameters of queues
-            that will be used to get incoming messages.
+        :param task_queues: parameters of queues that will be used to get incoming messages.
         :param dead_letter_queue: parameters of dead-letter queue.
         :param delay_queue: parameters of queue for simple delay implementation.
-        :param delayed_message_exchange_plugin: turn on or disable
-            delayed-message-exchange rabbitmq plugin.
-        :param delayed_message_exchange: parameters of exchange
-            that used to send messages with delay.
+        :param delayed_message_exchange_plugin: turn on or disable delayed-message-exchange rabbitmq plugin.
+        :param delayed_message_exchange: parameters of exchange that used to send messages with delay.
         :param label_for_routing: label name to use for routing key selection.
         :param label_for_priority: label name to use for message priority.
-        :param connection_kwargs: additional keyword arguments,
-            for connect_robust method of aio-pika.
+        :param connection_kwargs: additional keyword arguments, for connect_robust method of aio-pika.
         """
         super().__init__(result_backend, task_id_generator)
 
@@ -241,10 +236,8 @@ async def _declare_queues(
         """
         Declare all queues.
 
-        It's useful since aio-pika have automatic
-        recover mechanism, which works only if
-        the queue, you're going to listen was
-        declared by aio-pika.
+        It's useful since aio-pika have automatic recover mechanism, which works only if the queue, you're going to
+        listen was declared by aio-pika.
 
         :param channel: channel to used for declaration.
         :return: list of declared queues and their consumer arguments.
@@ -323,8 +316,7 @@ def with_queue(self, queue: Queue) -> Self:
         """
         Add new queue to the broker.
 
-        This method should be called before startup,
-        otherwise the new queue won't be declared and bound to exchange.
+        This method should be called before startup, otherwise the new queue won't be declared and bound to exchange.
 
         :param queue: queue to add.
         :return: self.
@@ -346,12 +338,8 @@ async def kick(self, message: BrokerMessage) -> None:
         """
         Send message to the exchange.
 
-        This function constructs rmq message
-        and sends it.
-
-        The message has task_id and task_name and labels
-        in headers. And message's routing key is the same
-        as the task_name.
+        This function constructs rmq message and sends it. The message has task_id and task_name and labels in headers.
+        And message's routing key is the same as the task_name.
 
         :raises NoStartupError: if startup wasn't called.
         :raises IncorrectRoutingKeyError: if routing key is incorrect.
@@ -420,8 +408,7 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]:
         """
         Listen to queue.
 
-        This function listens to queue and
-        yields every new message.
+        This function listens to queue and yields every new message.
 
         :raises NoStartupError: if startup wasn't called.
         :yields: parsed broker message.
@@ -446,17 +433,14 @@ async def body(
                 # Suppress errors during iterator cleanup if channel is being closed
                 logger.info("Queue iterator closed during shutdown")
 
-        combine = stream.merge(
+        async for message in merge_async_iterables(
             *[
                 body(queue, consumer_args)
                 for queue, consumer_args in queue_with_consumer_args_list
                 if not self._delay_queue or queue.name != self._delay_queue.name
             ],
-        )
-
-        async with combine.stream() as streamer:
-            async for message in streamer:
-                yield message
+        ):
+            yield message
 
     async def shutdown(self) -> None:
         """Close all connections on shutdown."""

taskiq_aio_pika/utils.py

@@ -0,0 +1,47 @@
+import asyncio
+from collections.abc import AsyncGenerator
+from typing import Any, TypeVar
+
+from typing_extensions import Sentinel
+
+_T = TypeVar("_T")
+_SENTINEL = Sentinel("_SENTINEL")
+
+
+async def merge_async_iterables(
+    *iterables: AsyncGenerator[_T, None],
+) -> AsyncGenerator[_T, None]:
+    """
+    Merge multiple async generators into a single one.
+
+    Items are yielded as soon as they're produced by any of the source generators, in the order they arrive.
+
+    :param iterables: async generators to merge.
+    :yields: items produced by any of the source generators.
+    """
+    queue: asyncio.Queue[Any] = asyncio.Queue()
+
+    async def _pump(iterable: AsyncGenerator[_T, None]) -> None:
+        try:
+            async for item in iterable:
+                await queue.put(item)
+        except BaseException as exc:
+            await queue.put(exc)
+        else:
+            await queue.put(_SENTINEL)
+
+    tasks = [asyncio.ensure_future(_pump(iterable)) for iterable in iterables]
+    remaining = len(tasks)
+    try:
+        while remaining:
+            item = await queue.get()
+            if item is _SENTINEL:
+                remaining -= 1
+            elif isinstance(item, BaseException):
+                raise item
+            else:
+                yield item
+    finally:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)

tests/test_utils.py

@@ -0,0 +1,96 @@
+import asyncio
+from collections.abc import AsyncGenerator
+
+import pytest
+
+from taskiq_aio_pika.utils import merge_async_iterables
+
+
+async def _gen(*items: int, delay: float = 0.0) -> AsyncGenerator[int, None]:
+    for item in items:
+        if delay:
+            await asyncio.sleep(delay)
+        yield item
+
+
+async def test_when_multiple_generators_passed__then_all_items_are_yielded() -> None:
+    result = [
+        item
+        async for item in merge_async_iterables(
+            _gen(1, 2, 3),
+            _gen(4, 5),
+        )
+    ]
+
+    assert sorted(result) == [1, 2, 3, 4, 5]
+
+
+async def test_when_single_generator_passed__then_its_items_are_yielded_in_order() -> (
+    None
+):
+    result = [item async for item in merge_async_iterables(_gen(1, 2, 3))]
+
+    assert result == [1, 2, 3]
+
+
+async def test_when_no_generators_passed__then_nothing_is_yielded() -> None:
+    result = [item async for item in merge_async_iterables()]
+
+    assert result == []
+
+
+async def test_when_one_generator_raises__then_exception_is_propagated() -> None:
+    async def failing_gen() -> AsyncGenerator[int, None]:
+        yield 1
+        raise ValueError("boom")
+
+    with pytest.raises(ValueError, match="boom"):
+        async for _ in merge_async_iterables(failing_gen(), _gen(2, delay=1)):
+            pass
+
+
+async def test_when_one_generator_raises__then_other_generators_are_cancelled() -> None:
+    other_was_cancelled = False
+
+    async def slow_gen() -> AsyncGenerator[int, None]:
+        nonlocal other_was_cancelled
+        try:
+            await asyncio.sleep(10)
+            yield 1
+        except asyncio.CancelledError:
+            other_was_cancelled = True
+            raise
+
+    async def failing_gen() -> AsyncGenerator[int, None]:
+        yield 1
+        raise ValueError("boom")
+
+    with pytest.raises(ValueError, match="boom"):
+        async for _ in merge_async_iterables(failing_gen(), slow_gen()):
+            pass
+
+    assert other_was_cancelled
+
+
+async def test_when_consumer_stops_early__then_remaining_generators_are_cancelled() -> (
+    None
+):
+    other_was_cancelled = False
+
+    async def slow_gen() -> AsyncGenerator[int, None]:
+        nonlocal other_was_cancelled
+        try:
+            await asyncio.sleep(10)
+            yield 1
+        except asyncio.CancelledError:
+            other_was_cancelled = True
+            raise
+
+    merged = merge_async_iterables(_gen(1), slow_gen())
+    try:
+        async for _ in merged:
+            break
+    finally:
+        await merged.aclose()
+
+    assert other_was_cancelled