asyncio-extensions 0.2.0__tar.gz → 0.3.0__tar.gz

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/.github/workflows/release.yml

@@ -13,7 +13,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
         with:
           persist-credentials: false
 

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/.github/workflows/tests.yml

@@ -19,7 +19,7 @@ jobs:
     strategy:
       fail-fast: false
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
 
       - name: Set up Python 3.13
         uses: actions/setup-python@v6
@@ -32,23 +32,27 @@ jobs:
       - name: Run static analysis
         run: hatch fmt --check
 
-      - name: Run type checking
+      - name: Type checking with MyPy
         run: hatch run mypy src/asyncio_extensions
 
+      - name: Type checking with Pyrefly
+        run: hatch run pyrefly check src/asyncio_extensions
+
   tests:
     name: Run tests on ${{ matrix.python-version }}
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.11", "3.12", "3.13", "3.14"]
+        python-version: ["3.11", "3.12", "3.13", "3.14", "3.15"]
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v7
 
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
+          allow-prereleases: true
 
       - name: Setup Hatch
         run: pipx  install hatch
@@ -58,13 +62,13 @@ jobs:
           hatch test -i py=${{ matrix.python-version }} --cover --junitxml=junit.xml
 
       - name: Upload results to Codecov
-        uses: codecov/codecov-action@v6
+        uses: codecov/codecov-action@v7
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
 
       - name: Upload test results to Codecov
         if: ${{ !cancelled() }}
-        uses: codecov/codecov-action@v6
+        uses: codecov/codecov-action@v7
         with:
           report_type: test_results
           token: ${{ secrets.CODECOV_TOKEN }}

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/commitizen-tools/commitizen
-    rev: v4.15.1
+    rev: v4.16.4
     hooks:
       - id: commitizen
       - id: commitizen-branch

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/AGENTS.md

@@ -12,7 +12,7 @@ All implementation lives under `src/asyncio_extensions/`. Every submodule is **p
 
 | Module | Theme |
 |---|---|
-| `_compat.py` | Version compatibility shims (Python 3.11–3.14). No new features here. |
+| `_compat.py` | Version compatibility shims (Python 3.11–3.15). No new features here. |
 | `_sync.py` | Bridging sync and async code. |
 | `_scheduling.py` | Event-loop control utilities. |
 | `_iterators.py` | Async iteration utilities. |
@@ -55,7 +55,7 @@ from asyncio_extensions import checkpoint  # helper from another module
 
 Maintain the same order as the source module: tests for each symbol should appear in the order the symbols are defined, and within a symbol's tests, individual cases should follow the order of branches and conditions in the implementation.
 
-The test suite runs against Python 3.11, 3.12, 3.13, and 3.14 in CI. When writing version-specific tests, use `@pytest.mark.skipif(sys.version_info < (3, X), ...)`.
+The test suite runs against Python 3.11, 3.12, 3.13, 3.14 and 3.15 in CI. When writing version-specific tests, use `@pytest.mark.skipif(sys.version_info < (3, X), ...)`.
 
 ## Compatibility
 

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/CHANGELOG.md

@@ -8,6 +8,18 @@ This project adheres to both [Semantic Versioning](https://semver.org/spec/v2.0.
 
 ## [UNRELEASED]
 
+### Added
+
+- `drain` coroutine to consume and discard remaining items from an async iterator or async iterable.
+- `merge_streams` managed-stream helper to merge multiple `ManagedStream` sources into a single interleaved stream.
+
+### Changed
+
+- Type validation now includes Pyrefly checks for enhanced type safety.
+- Enabled mypy strict mode to enforce stricter type checking standards.
+- `STOP` is now a proper sentinel instance (via `typing_extensions.sentinel` on Python < 3.15) instead of a class used as a sentinel. Behaviour is unchanged; only the type changes.
+- Use built-in TaskGroup on Python 3.15+
+
 ## [0.2.0] - 2026-05-15
 
 ### Added

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: asyncio-extensions
-Version: 0.2.0
+Version: 0.3.0
 Summary: Some (maybe) useful extensions for asyncio
 Project-URL: Changelog, https://github.com/hartungstenio/asyncio-extensions/blob/main/CHANGELOG.md
 Project-URL: Documentation, https://github.com/hartungstenio/asyncio-extensions#readme
@@ -19,12 +19,13 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Typing :: Typed
 Requires-Python: >=3.11
-Requires-Dist: typing-extensions; python_version < '3.13'
+Requires-Dist: typing-extensions>=4.16; python_version < '3.15'
 Description-Content-Type: text/markdown
 
 # asyncio-extensions
@@ -269,6 +270,25 @@ async def source():
 await fill_queue(source(), queue)
 ```
 
+### drain
+
+The `drain` coroutine consumes and discards the remaining items from an async iterator or async iterable. This is useful when you want to stop processing early but must ensure the source is exhausted so producers can finish and resources can be released.
+
+Example — process only the first N items, then discard the rest so the producer can complete:
+
+```python
+from asyncio_extensions import drain
+
+async def iter_items():
+    ...  # yields items that need to process
+
+async for idx, item in enumerate(iter_items()):
+    if idx >= 5:
+        break
+    print(item)
+# ensure the remaining items are consumed and the producer can finish
+await drain(stream)
+```
 ### merge_iterables
 
 The `merge_iterables` async context manager merges multiple sync or async iterables into a single interleaved stream, feeding all sources into a shared queue concurrently.
@@ -306,6 +326,29 @@ def my_stream() -> ManagedStream[int]:
     ...
 ```
 
+### merge_streams
+
+`merge_streams` is the managed-stream equivalent of `merge_iterables`. It accepts one or more `ManagedStream` objects (context managers that yield `AsyncIterator`s) and returns a single managed stream that interleaves items from all sources. Use it when each source requires context-managed setup/cleanup (for example, generators wrapped with `safe_gen` or `asynccontextmanager`-based streams).
+
+Example — merging two managed streams produced by `safe_gen`:
+
+```python
+from asyncio_extensions import merge_streams, safe_gen, checkpoint
+
+@safe_gen
+async def source(name: str):
+    for i in range(3):
+        yield f"{name}-{i}"
+        await checkpoint()
+
+async def main():
+    async with merge_streams(source("a"), source("b")) as stream:
+        async for item in stream:
+            print(item)
+```
+
+Each input managed stream is guaranteed to be closed when the merged context exits, ensuring proper cleanup of resources.
+
 ### safe_gen
 
 The `safe_gen` decorator converts an async generator function into a context manager, enforcing correct handling of early exits. A plain async generator abandoned before exhaustion leaks resources and keeps running indefinitely — the caller has no way to know it needs to call `aclose()`. By returning a context manager instead, `safe_gen` makes cleanup syntactically mandatory: callers must use `async with`, which guarantees `aclose()` is called on exit regardless of how iteration ends. It also suppresses `GeneratorExit` raised inside an exception group, so it composes safely with `TaskGroup`.

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/README.md

@@ -240,6 +240,25 @@ async def source():
 await fill_queue(source(), queue)
 ```
 
+### drain
+
+The `drain` coroutine consumes and discards the remaining items from an async iterator or async iterable. This is useful when you want to stop processing early but must ensure the source is exhausted so producers can finish and resources can be released.
+
+Example — process only the first N items, then discard the rest so the producer can complete:
+
+```python
+from asyncio_extensions import drain
+
+async def iter_items():
+    ...  # yields items that need to process
+
+async for idx, item in enumerate(iter_items()):
+    if idx >= 5:
+        break
+    print(item)
+# ensure the remaining items are consumed and the producer can finish
+await drain(stream)
+```
 ### merge_iterables
 
 The `merge_iterables` async context manager merges multiple sync or async iterables into a single interleaved stream, feeding all sources into a shared queue concurrently.
@@ -277,6 +296,29 @@ def my_stream() -> ManagedStream[int]:
     ...
 ```
 
+### merge_streams
+
+`merge_streams` is the managed-stream equivalent of `merge_iterables`. It accepts one or more `ManagedStream` objects (context managers that yield `AsyncIterator`s) and returns a single managed stream that interleaves items from all sources. Use it when each source requires context-managed setup/cleanup (for example, generators wrapped with `safe_gen` or `asynccontextmanager`-based streams).
+
+Example — merging two managed streams produced by `safe_gen`:
+
+```python
+from asyncio_extensions import merge_streams, safe_gen, checkpoint
+
+@safe_gen
+async def source(name: str):
+    for i in range(3):
+        yield f"{name}-{i}"
+        await checkpoint()
+
+async def main():
+    async with merge_streams(source("a"), source("b")) as stream:
+        async for item in stream:
+            print(item)
+```
+
+Each input managed stream is guaranteed to be closed when the merged context exits, ensuring proper cleanup of resources.
+
 ### safe_gen
 
 The `safe_gen` decorator converts an async generator function into a context manager, enforcing correct handling of early exits. A plain async generator abandoned before exhaustion leaks resources and keeps running indefinitely — the caller has no way to know it needs to call `aclose()`. By returning a context manager instead, `safe_gen` makes cleanup syntactically mandatory: callers must use `async with`, which guarantees `aclose()` is called on exit regardless of how iteration ends. It also suppresses `GeneratorExit` raised inside an exception group, so it composes safely with `TaskGroup`.

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/pyproject.toml

@@ -18,13 +18,14 @@ classifiers = [
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
   "Programming Language :: Python :: 3.14",
+  "Programming Language :: Python :: 3.15",
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
   "Topic :: Software Development :: Libraries",
   "Typing :: Typed",
 ]
 dependencies = [
-  "typing-extensions; python_version < '3.13'",
+  "typing-extensions (>= 4.16); python_version < '3.15'",
 ]
 dynamic = ["version"]
 
@@ -40,6 +41,7 @@ dev = [
   "faker",
   "mypy",
   "pre-commit",
+  "pyrefly",
   "pytest",
   "pytest-asyncio",
   "pytest-cov",
@@ -87,7 +89,7 @@ extra-dependencies = [
 ]
 
 [[tool.hatch.envs.hatch-test.matrix]]
-python = ["3.11", "3.12", "3.13", "3.14"]
+python = ["3.11", "3.12", "3.13", "3.14", "3.15"]
 
 [tool.hatch.envs.hatch-test.scripts]
 run = "pytest{env:HATCH_TEST_ARGS:} {args}"
@@ -95,6 +97,14 @@ run-cov = "coverage run -m pytest{env:HATCH_TEST_ARGS:} {args}"
 cov-combine = "coverage combine"
 cov-report = "coverage xml"
 
+[tool.mypy]
+strict = true
+enable_error_code = ["deprecated"]
+
+[tool.pyrefly]
+preset = "strict"
+min-severity = "ignore"
+
 [tool.pytest]
 asyncio_default_fixture_loop_scope = "function"
 

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/src/asyncio_extensions/__init__.py

@@ -4,10 +4,12 @@ from ._compat import CreateTaskParams as CreateTaskParams
 from ._compat import iscoroutinefunction as iscoroutinefunction
 from ._iterators import STOP as STOP
 from ._iterators import ManagedStream as ManagedStream
+from ._iterators import drain as drain
 from ._iterators import fill_queue as fill_queue
 from ._iterators import flatten_stream as flatten_stream
 from ._iterators import iterate_queue as iterate_queue
 from ._iterators import merge_iterables as merge_iterables
+from ._iterators import merge_streams as merge_streams
 from ._iterators import safe_gen as safe_gen
 from ._scheduling import checkpoint as checkpoint
 from ._scheduling import heartbeat as heartbeat

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/src/asyncio_extensions/_compat.py

@@ -45,6 +45,12 @@ if sys.version_info >= (3, 14):
 else:
     CreateTaskParams = _CreateTaskParams
 
+if sys.version_info >= (3, 15):
+    import builtins
+
+    sentinel = builtins.sentinel
+else:
+    from typing_extensions import sentinel  # type: ignore[attr-defined]
 
 __all__ = [
     "CreateTaskParams",
@@ -53,4 +59,5 @@ __all__ = [
     "iscoroutinefunction",
     "markcoroutinefunction",
     "override",
+    "sentinel",
 ]

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/src/asyncio_extensions/_iterators.py

@@ -3,11 +3,11 @@
 import asyncio
 import sys
 from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Callable, Iterable
-from contextlib import AbstractAsyncContextManager, aclosing, asynccontextmanager
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, aclosing, asynccontextmanager
 from functools import wraps
-from typing import Any, ParamSpec, TypeAlias, TypeVar
+from typing import ParamSpec, TypeAlias, TypeVar
 
-from ._compat import QueueShutDown
+from ._compat import QueueShutDown, sentinel
 from ._sync import asyncify_iterable
 from ._task_groups import TaskGroup
 
@@ -25,30 +25,26 @@ that return a context-managed async stream::
 """
 
 if sys.version_info >= (3, 13):
-    from warnings import deprecated
+    STOP = sentinel("STOP")
+    """Sentinel that signals the end of an :func:`iterate_queue` stream.
 
-    @deprecated("Prefer using Queue.shutdown instead.")
-    class STOP:
-        """Sentinel that signals the end of an :func:`iterate_queue` stream.
+    Put this class in the queue to stop iteration::
 
-        Put this class in the queue to stop iteration::
+        await queue.put(STOP)
 
-            await queue.put(STOP)
-
-        Deprecated:
-            Prefer :meth:`asyncio.Queue.shutdown` on Python 3.13+.
-        """
+    Deprecated:
+        Prefer :meth:`asyncio.Queue.shutdown` on Python 3.13+.
+    """
 
 
 else:
+    STOP = sentinel("STOP")
+    """Sentinel that signals the end of an :func:`iterate_queue` stream.
 
-    class STOP:
-        """Sentinel that signals the end of an :func:`iterate_queue` stream.
+    Put this class in the queue to stop iteration::
 
-        Put this class in the queue to stop iteration::
-
-            await queue.put(STOP)
-        """
+        await queue.put(STOP)
+    """
 
 
 async def iterate_queue(queue: asyncio.Queue[T]) -> AsyncGenerator[T]:
@@ -100,10 +96,26 @@ async def fill_queue(itr: AsyncIterable[T] | Iterable[T], queue: asyncio.Queue[T
         await queue.put(it)
 
 
+async def drain(itr: AsyncIterable[T] | Iterable[T]) -> None:
+    """Consume the remaining items from *itr*.
+
+    Accepts both sync and async iterables.
+
+    Args:
+        itr: The iterable (sync or async) to consume.
+
+    Example::
+
+        await drain(range(10))
+    """
+    async for _ in asyncify_iterable(itr):
+        pass
+
+
 @asynccontextmanager
 async def merge_iterables(
     *itrs: AsyncIterable[T] | Iterable[T],
-) -> AsyncIterator[AsyncGenerator[T]]:
+) -> AsyncGenerator[AsyncGenerator[T]]:
     """Merge multiple iterables into a single async stream.
 
     Feeds all *itrs* into a shared queue concurrently and yields a single
@@ -133,7 +145,7 @@ async def merge_iterables(
 
         async def join() -> None:
             await asyncio.wait(tasks)
-            await queue.put(STOP)  # type: ignore[arg-type]
+            await queue.put(STOP)
 
         tg.create_task(join())
 
@@ -143,7 +155,32 @@ async def merge_iterables(
             tg.cancel()
 
 
-def safe_gen(fn: Callable[P, AsyncGenerator[T, Any]]) -> Callable[P, ManagedStream[T]]:
+@asynccontextmanager
+async def merge_streams(*streams: ManagedStream[T]) -> AsyncGenerator[AsyncGenerator[T]]:
+    """Merge multiple managed streams into a single async stream.
+
+    Similar to :func:`merge_iterables`, but work with managed streams.
+
+    Args:
+        *streams: Any number of managed streams to merge.
+
+    Yields:
+        An async generator producing items from all *itrs* interleaved.
+
+    Example::
+
+        async with merge_streams(source_a, source_b) as stream:
+            async for item in stream:
+                process(item)
+    """
+    async with AsyncExitStack() as exit_stack:
+        itrs = [await exit_stack.enter_async_context(stream) for stream in streams]
+        merged = await exit_stack.enter_async_context(merge_iterables(*itrs))
+
+        yield merged
+
+
+def safe_gen(fn: Callable[P, AsyncGenerator[T]]) -> Callable[P, ManagedStream[T]]:
     """Wrap an async generator function so it is always closed on exit.
 
     Async generators must be explicitly closed when iteration is abandoned early — otherwise
@@ -169,7 +206,7 @@ def safe_gen(fn: Callable[P, AsyncGenerator[T, Any]]) -> Callable[P, ManagedStre
 
     @asynccontextmanager
     @wraps(fn)
-    async def decorator(*args: P.args, **kwargs: P.kwargs) -> AsyncIterator[AsyncGenerator[T]]:
+    async def decorator(*args: P.args, **kwargs: P.kwargs) -> AsyncGenerator[AsyncGenerator[T]]:
         try:
             async with aclosing(fn(*args, **kwargs)) as agen:
                 yield agen

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/src/asyncio_extensions/_scheduling.py

@@ -2,7 +2,7 @@
 
 import asyncio
 from collections.abc import Awaitable, Callable, Generator
-from typing import Any, Never, ParamSpec, TypeVar
+from typing import Never, ParamSpec, TypeVar
 
 T = TypeVar("T")
 P = ParamSpec("P")
@@ -11,7 +11,7 @@ P = ParamSpec("P")
 class _YieldToEventLoop:
     """Helper class to give control back to the event loop."""
 
-    def __await__(self) -> Generator[None, Any, None]:
+    def __await__(self) -> Generator[None]:
         """Yield control to the event loop."""
         yield
 

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/src/asyncio_extensions/_sync.py

@@ -53,7 +53,7 @@ def asyncify(func: Callable[_P, _R] | Callable[_P, Awaitable[_R]], /) -> Callabl
     return wrapper
 
 
-def markcoroutinefunction(f: Callable[_P, _R]) -> Callable[_P, Coroutine[Any, Any, _R]]:
+def markcoroutinefunction(f: Callable[_P, _R]) -> Callable[_P, Coroutine[None, None, _R]]:
     """Mark *f* as a coroutine function without making it async.
 
     After marking, :func:`asyncio.iscoroutinefunction` returns ``True``

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/src/asyncio_extensions/_task_groups.py

@@ -1,9 +1,10 @@
 """Extensions to :mod:`asyncio` task groups."""
 
 import asyncio
+import sys
 from collections.abc import Coroutine
 from types import TracebackType
-from typing import Any, Never, TypeVar, Unpack
+from typing import Never, TypeVar, Unpack
 
 from ._compat import CreateTaskParams, override
 
@@ -24,23 +25,28 @@ async def force_terminate_task_group() -> Never:
     raise TerminateTaskGroup
 
 
-class TaskGroup(asyncio.TaskGroup):
-    """An :class:`asyncio.TaskGroup` extended with a :meth:`cancel` method."""
+if sys.version_info >= (3, 15):
+    from asyncio import TaskGroup
+else:
 
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> None:
-        try:
-            await super().__aexit__(exc_type, exc_value, traceback)
-        except* TerminateTaskGroup:
-            pass
+    class TaskGroup(asyncio.TaskGroup):
+        """An :class:`asyncio.TaskGroup` extended with a :meth:`cancel` method."""
+
+        @override
+        async def __aexit__(
+            self,
+            et: type[BaseException] | None,
+            exc: BaseException | None,
+            tb: TracebackType | None,
+        ) -> None:
+            try:
+                await super().__aexit__(et, exc, tb)
+            except* TerminateTaskGroup:
+                pass
 
-    def cancel(self) -> None:
-        """Schedule cancellation of all remaining tasks in the group."""
-        self.create_task(force_terminate_task_group())
+        def cancel(self) -> None:
+            """Schedule cancellation of all remaining tasks in the group."""
+            self.create_task(force_terminate_task_group())
 
 
 class LimitedTaskGroup(TaskGroup):
@@ -54,14 +60,14 @@ class LimitedTaskGroup(TaskGroup):
         super().__init__()
         self._semaphore = asyncio.Semaphore(max_concurrent)
 
-    async def _task_wrapper(self, coro: Coroutine[Any, Any, T]) -> T:
+    async def _task_wrapper(self, coro: Coroutine[None, None, T]) -> T:
         async with self._semaphore:
             return await coro
 
     @override
     def create_task(
         self,
-        coro: Coroutine[Any, Any, T],
+        coro: Coroutine[None, None, T],
         **kwargs: Unpack[CreateTaskParams],
     ) -> asyncio.Task[T]:
         """Schedule *coro* as a task, subject to the concurrency limit.

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/tests/test_iterators.py

@@ -1,20 +1,31 @@
 import asyncio
 import sys
 from collections.abc import AsyncGenerator
+from contextlib import aclosing
 from itertools import count
 
 import pytest
 
 from asyncio_extensions import checkpoint
-from asyncio_extensions._iterators import STOP, fill_queue, flatten_stream, iterate_queue, merge_iterables, safe_gen
+from asyncio_extensions._iterators import (
+    STOP,
+    drain,
+    fill_queue,
+    flatten_stream,
+    iterate_queue,
+    merge_iterables,
+    merge_streams,
+    safe_gen,
+)
 
 pytestmark = pytest.mark.asyncio
 
 
-def drain(queue: asyncio.Queue[int]) -> list[int]:
+def _drain_queue(queue: asyncio.Queue[int]) -> list[int]:
     items = []
     while not queue.empty():
         items.append(queue.get_nowait())
+        queue.task_done()
     return items
 
 
@@ -88,7 +99,7 @@ async def test_fill_queue_sync_iterable_fills_queue() -> None:
 
     await fill_queue([1, 2, 3], queue)
 
-    assert drain(queue) == [1, 2, 3]
+    assert _drain_queue(queue) == [1, 2, 3]
 
 
 async def test_fill_queue_async_iterable_fills_queue() -> None:
@@ -100,7 +111,7 @@ async def test_fill_queue_async_iterable_fills_queue() -> None:
 
     await fill_queue(source(), queue)
 
-    assert drain(queue) == [4, 5, 6]
+    assert _drain_queue(queue) == [4, 5, 6]
 
 
 async def test_fill_queue_empty_iterable_queue_remains_empty() -> None:
@@ -123,7 +134,32 @@ async def test_fill_queue_full_queue_blocks_until_space_available() -> None:
 
         queue.get_nowait()
 
-    assert drain(queue) == [1]
+    assert _drain_queue(queue) == [1]
+
+
+# drain
+
+
+async def test_drain_sync_iterator() -> None:
+    itr = iter([1, 2, 3])
+
+    await drain(itr)
+
+    with pytest.raises(StopIteration):
+        next(itr)
+
+
+async def test_drain_async_iterator() -> None:
+    async def gen() -> AsyncGenerator[int]:
+        for i in range(3):
+            yield i
+
+    itr = aiter(gen())
+
+    await drain(itr)
+
+    with pytest.raises(StopAsyncIteration):
+        await anext(itr)
 
 
 # merge_iterables
@@ -174,6 +210,58 @@ async def test_merge_iterables_early_exit_cancels_background_tasks() -> None:
     assert len(asyncio.all_tasks()) == initial_tasks
 
 
+# merge_streams
+
+
+async def test_merge_streams_yields_all_items() -> None:
+    async def source(values: list[int]) -> AsyncGenerator[int]:
+        for v in values:
+            yield v
+
+    async with merge_streams(aclosing(source([10, 20])), aclosing(source([30, 40]))) as stream:
+        results = [item async for item in stream]
+
+    assert sorted(results) == [10, 20, 30, 40]
+
+
+async def test_merge_streams_single_source_yields_all_items() -> None:
+    async def source(values: list[int]) -> AsyncGenerator[int]:
+        for v in values:
+            yield v
+
+    async with merge_streams(aclosing(source([1, 2, 3]))) as stream:
+        results = [item async for item in stream]
+
+    assert results == [1, 2, 3]
+
+
+async def test_merge_streams_empty_sources_yields_nothing() -> None:
+    async def source() -> AsyncGenerator[int]:
+        if False:
+            yield
+
+    async with merge_streams(aclosing(source())) as stream:
+        results = [item async for item in stream]
+
+    assert results == []
+
+
+async def test_merge_streams_early_exit_cancels_background_tasks() -> None:
+    initial_tasks = len(asyncio.all_tasks())
+
+    async def source(values: list[int]) -> AsyncGenerator[int]:
+        for v in values:
+            yield v
+
+    async with merge_streams(aclosing(source([10, 20])), aclosing(source([30, 40]))) as stream:
+        assert len(asyncio.all_tasks()) > initial_tasks
+
+        async for _ in stream:
+            break
+
+    assert len(asyncio.all_tasks()) == initial_tasks
+
+
 async def test_safe_gen_full_iteration() -> None:
     @safe_gen
     async def gen(items: int) -> AsyncGenerator[int]:

{asyncio_extensions-0.2.0 → asyncio_extensions-0.3.0}/tests/test_task_groups.py

@@ -24,6 +24,7 @@ async def test_force_terminate_task_group() -> None:
             tg.create_task(force_terminate_task_group())
 
 
+@pytest.mark.skipif(sys.version_info >= (3, 15), reason="Skipping testing built-in TaskGroup")
 class TestTaskGroup:
     async def test_runs_tasks(self) -> None:
         async with TaskGroup() as tg: