Preserve custom acontainer factory return types

_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}

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):