diff --git a/_python_utils_tests/test_aio.py b/_python_utils_tests/test_aio.py index f9461fe..7125844 100644 --- a/_python_utils_tests/test_aio.py +++ b/_python_utils_tests/test_aio.py @@ -114,3 +114,25 @@ async def empty_gen() -> types.AsyncIterable[types.Tuple[int, int]]: assert await aio.adict(async_gen()) == {1: 2, 3: 4, 5: 6} assert await aio.adict(empty_gen) == {} assert await aio.adict(empty_gen()) == {} + + +@pytest.mark.asyncio +async def test_acontainer_custom_container_types() -> None: + """Accept custom container factories without losing their return type.""" + from collections.abc import AsyncIterator, Iterable + + from typing_extensions import assert_type + + async def numbers() -> AsyncIterator[int]: + yield 1 + yield 2 + + def summarize(items: Iterable[int]) -> dict[str, int]: + return {'total': sum(items)} + + frozen = await aio.acontainer(numbers(), frozenset) + assert_type(frozen, frozenset[int]) + assert frozen == frozenset({1, 2}) + summary = await aio.acontainer(numbers, summarize) + assert_type(summary, dict[str, int]) + assert summary == {'total': 3} diff --git a/python_utils/aio.py b/python_utils/aio.py index 2b0854d..2cab6df 100644 --- a/python_utils/aio.py +++ b/python_utils/aio.py @@ -18,6 +18,8 @@ _K = typing.TypeVar('_K') #: Value type when collecting an async iterable of pairs into a dict. _V = typing.TypeVar('_V') +#: Concrete collection returned by a custom container factory. +_C = typing.TypeVar('_C', bound=collections.abc.Collection[typing.Any]) async def acount( @@ -83,6 +85,15 @@ async def acontainer( """Overload: collect the async iterable into a ``set``.""" +@typing.overload +async def acontainer( + iterable: collections.abc.AsyncIterable[_T] + | collections.abc.Callable[..., collections.abc.AsyncIterable[_T]], + container: collections.abc.Callable[[collections.abc.Iterable[_T]], _C], +) -> _C: + """Overload: preserve the result type of a custom container factory.""" + + async def acontainer( iterable: collections.abc.AsyncIterable[_T] | collections.abc.Callable[..., collections.abc.AsyncIterable[_T]], @@ -96,6 +107,8 @@ async def acontainer( So instead of doing `[item async for item in iterable]` you can do `await acontainer(iterable)`. + ``container`` may be any factory that accepts the collected items, such + as ``frozenset`` or a custom function. Its concrete return type is kept. """ iterable_: collections.abc.AsyncIterable[_T] if callable(iterable):