From 88dfa98ca8830428801e839f09c1f9ad3eeb6be3 Mon Sep 17 00:00:00 2001 From: Dima Anfimov Date: Wed, 16 Sep 2026 14:33:00 +0200 Subject: [PATCH 1/2] feat: remove aiostream as a dependency --- LICENSE | 2 +- Makefile | 8 ++++ pyproject.toml | 4 +- taskiq_aio_pika/broker.py | 53 +++++++++------------ taskiq_aio_pika/utils.py | 51 +++++++++++++++++++++ tests/test_utils.py | 96 +++++++++++++++++++++++++++++++++++++++ uv.lock | 18 ++------ 7 files changed, 182 insertions(+), 50 deletions(-) create mode 100644 taskiq_aio_pika/utils.py create mode 100644 tests/test_utils.py diff --git a/LICENSE b/LICENSE index b44346c..aad3c50 100644 --- a/LICENSE +++ b/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 diff --git a/Makefile b/Makefile index 01a59e0..1bab2ec 100644 --- a/Makefile +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index cff19da..c116770 100644 --- a/pyproject.toml +++ b/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 ; python_version < '3.15'", ] [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/*" = [ diff --git a/taskiq_aio_pika/broker.py b/taskiq_aio_pika/broker.py index 0cc9601..cbd2a95 100644 --- a/taskiq_aio_pika/broker.py +++ b/taskiq_aio_pika/broker.py @@ -1,4 +1,5 @@ import asyncio +import sys from collections.abc import AsyncGenerator, Callable from datetime import timedelta from logging import getLogger @@ -7,10 +8,13 @@ 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 + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self from taskiq_aio_pika.exceptions import ( ExchangeNotDeclaredError, @@ -20,6 +24,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 +74,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 +241,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 +321,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 +343,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 +413,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 +438,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.""" diff --git a/taskiq_aio_pika/utils.py b/taskiq_aio_pika/utils.py new file mode 100644 index 0000000..654422d --- /dev/null +++ b/taskiq_aio_pika/utils.py @@ -0,0 +1,51 @@ +import asyncio +import sys +from collections.abc import AsyncGenerator +from typing import Any, TypeVar + +if sys.version_info >= (3, 15): + from typing import Sentinel +else: + 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) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..d19073e --- /dev/null +++ b/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 diff --git a/uv.lock b/uv.lock index 445bb47..484a88f 100644 --- a/uv.lock +++ b/uv.lock @@ -171,18 +171,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] -[[package]] -name = "aiostream" -version = "0.7.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/8b/65/b9b69695702b76a878c9879f2ee80cefce75bc5cb864fc100460bc1c5380/aiostream-0.7.1.tar.gz", hash = "sha256:272aaa0d8f83beb906f5aa9022bb59046bb7a103fa3770f807c31f918595acf6", size = 44059, upload-time = "2025-10-13T20:02:06.961Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/52/a0/d7c6ca304140f3f49987d710e15bc164248924a35d8cdfac2f6e87fca041/aiostream-0.7.1-py3-none-any.whl", hash = "sha256:ea8739e9158ee6a606b3feedf3762721c3507344e540d09a10984c5e88a13b37", size = 41416, upload-time = "2025-10-13T20:02:05.535Z" }, -] - [[package]] name = "annotated-types" version = "0.7.0" @@ -394,7 +382,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" } wheels = [ @@ -1196,8 +1184,8 @@ version = "0.6.0" source = { editable = "." } dependencies = [ { name = "aio-pika" }, - { name = "aiostream" }, { name = "taskiq" }, + { name = "typing-extensions", marker = "python_full_version < '3.15'" }, ] [package.dev-dependencies] @@ -1234,8 +1222,8 @@ typecheck = [ [package.metadata] requires-dist = [ { name = "aio-pika", specifier = ">=9.0.0" }, - { name = "aiostream", specifier = ">=0.7.1" }, { name = "taskiq", specifier = ">=0.12.0,<1" }, + { name = "typing-extensions", marker = "python_full_version < '3.15'", specifier = ">=4.14.0" }, ] [package.metadata.requires-dev] From a1c8e9990a9e727783791802f4f8842183b28872 Mon Sep 17 00:00:00 2001 From: Dima Anfimov Date: Wed, 16 Sep 2026 14:40:11 +0200 Subject: [PATCH 2/2] fix: use typing_extension unconditionally --- pyproject.toml | 2 +- taskiq_aio_pika/broker.py | 7 +------ taskiq_aio_pika/utils.py | 6 +----- 3 files changed, 3 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c116770..0f7c802 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ requires-python = ">=3.10,<4" dependencies = [ "taskiq>=0.12.0,<1", "aio-pika>=9.0.0", - "typing-extensions>=4.14.0 ; python_version < '3.15'", + "typing-extensions>=4.14.0", ] [dependency-groups] diff --git a/taskiq_aio_pika/broker.py b/taskiq_aio_pika/broker.py index cbd2a95..764ccdc 100644 --- a/taskiq_aio_pika/broker.py +++ b/taskiq_aio_pika/broker.py @@ -1,5 +1,4 @@ import asyncio -import sys from collections.abc import AsyncGenerator, Callable from datetime import timedelta from logging import getLogger @@ -10,11 +9,7 @@ from aio_pika.abc import AbstractChannel, AbstractQueue, AbstractRobustConnection from pamqp.common import FieldTable from taskiq import AckableMessage, AsyncBroker, AsyncResultBackend, BrokerMessage - -if sys.version_info >= (3, 11): - from typing import Self -else: - from typing_extensions import Self +from typing_extensions import Self from taskiq_aio_pika.exceptions import ( ExchangeNotDeclaredError, diff --git a/taskiq_aio_pika/utils.py b/taskiq_aio_pika/utils.py index 654422d..a4cd710 100644 --- a/taskiq_aio_pika/utils.py +++ b/taskiq_aio_pika/utils.py @@ -1,12 +1,8 @@ import asyncio -import sys from collections.abc import AsyncGenerator from typing import Any, TypeVar -if sys.version_info >= (3, 15): - from typing import Sentinel -else: - from typing_extensions import Sentinel +from typing_extensions import Sentinel _T = TypeVar("_T") _SENTINEL = Sentinel("_SENTINEL")