Repository: maxfischer2781/asyncstdlib

Issue:
Add heapq / heapq.merge?
The heapq module has a candidate to benefit from an ``async`` iterator variant, namely [``heapq.merge``](https://docs.python.org/3/library/heapq.html#heapq.merge).

* Is this a useful helper? Would other ``heapq`` helpers benefit?
    * I'd say "yes". It covers infinite/lazy iterables better than ``sorted``.
    * The ``nlargest`` and ``nsmallest`` would benefit from low memory compared to ``sorted``+slicing.
    * The ``heapq`` itself would likely not benefit; an ``async`` key function can be used as ``heapq.heappush(heap, (await key(item), item))`` as for the sync case.
* Should this go to a single-member ``asyncstdlib.heapq`` module?
    * [BPO#38938](https://bugs.python.org/issue38938) suggest ``merge`` might just as well live in ``itertools``.
    * ``asynctools`` might serve as a temporary home.

See also https://github.com/python/cpython/pull/20550 for the (candidate for) the most performant implementation.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer) Report
> Merging [#75](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer) (cad7e3d) into [master](https://codecov.io/gh/maxfischer2781/asyncstdlib/commit/9bbf8388e0021700d6967317e0c1112b5bbae44e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer) (9bbf838) will **decrease** coverage by `0.08%`.
> The diff coverage is `97.18%`.

[![Impacted file tree graph](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75/graphs/tree.svg?width=650&height=150&src=pr&token=frT6svSwUb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer)](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer)

```diff
@@            Coverage Diff             @@
##           master      #75      +/-   ##
==========================================
- Coverage   99.20%   99.11%   -0.09%     
==========================================
  Files          17       19       +2     
  Lines        1628     1699      +71     
  Branches      220      258      +38     
==========================================
+ Hits         1615     1684      +69     
- Misses          9       11       +2     
  Partials        4        4              
```


| [Impacted Files](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer) | Coverage Δ | |
|---|---|---|
| [asyncstdlib/heapq.py](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer#diff-YXN5bmNzdGRsaWIvaGVhcHEucHk=) | `96.29% <96.29%> (ø)` | |
| [asyncstdlib/\_\_init\_\_.py](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer#diff-YXN5bmNzdGRsaWIvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [unittests/test\_heapq.py](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer#diff-dW5pdHRlc3RzL3Rlc3RfaGVhcHEucHk=) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer). Last update [9bbf838...cad7e3d](https://codecov.io/gh/maxfischer2781/asyncstdlib/pull/75?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Max+Fischer).


Patch:
diff --git a/asyncstdlib/__init__.py b/asyncstdlib/__init__.py
index 0fd50b1..049121f 100644
--- a/asyncstdlib/__init__.py
+++ b/asyncstdlib/__init__.py
@@ -34,6 +34,7 @@ from .itertools import (
     groupby,
 )
 from .asynctools import borrow, scoped_iter, await_each, apply, sync
+from .heapq import merge, nlargest, nsmallest
 
 __version__ = "3.10.2"
 
@@ -83,4 +84,8 @@ __all__ = [
     "await_each",
     "apply",
     "sync",
+    # heapq
+    "merge",
+    "nlargest",
+    "nsmallest",
 ]
diff --git a/asyncstdlib/heapq.py b/asyncstdlib/heapq.py
new file mode 100644
index 0000000..8996f4a
--- /dev/null
+++ b/asyncstdlib/heapq.py
@@ -0,0 +1,266 @@
+from typing import (
+    Generic,
+    AsyncIterator,
+    Tuple,
+    List,
+    Optional,
+    Callable,
+    Any,
+    overload,
+    Awaitable,
+)
+import heapq as _heapq
+
+from .builtins import enumerate as a_enumerate, zip as a_zip
+from ._core import aiter, awaitify, ScopedIter, borrow
+from ._typing import AnyIterable, LT, T, SupportsLT
+
+
+class _KeyIter(Generic[LT]):
+    __slots__ = ("head", "tail", "reverse", "head_key", "key")
+
+    @overload
+    def __init__(
+        self,
+        head: T,
+        tail: AsyncIterator[T],
+        reverse: bool,
+        head_key: LT,
+        key: Callable[[T], Awaitable[LT]],
+    ):
+        pass
+
+    @overload
+    def __init__(
+        self, head: LT, tail: AsyncIterator[LT], reverse: bool, head_key: LT, key: None
+    ):
+        pass
+
+    def __init__(
+        self,
+        head: Any,
+        tail: AsyncIterator[Any],
+        reverse: bool,
+        head_key: LT,
+        key: Any,
+    ):
+        self.head = head
+        self.head_key = head_key
+        self.tail = tail
+        self.key = key
+        self.reverse = reverse
+
+    @overload
+    @classmethod
+    def from_iters(
+        cls,
+        iterables: Tuple[AnyIterable[T], ...],
+        reverse: bool,
+        key: Callable[[T], Awaitable[LT]],
+    ) -> "AsyncIterator[_KeyIter[LT]]":
+        pass
+
+    @overload
+    @classmethod
+    def from_iters(
+        cls, iterables: Tuple[AnyIterable[LT], ...], reverse: bool, key: None
+    ) -> "AsyncIterator[_KeyIter[LT]]":
+        pass
+
+    @classmethod
+    async def from_iters(
+        cls,
+        iterables: Tuple[AnyIterable[Any], ...],
+        reverse: bool,
+        key: Optional[Callable[[Any], Any]],
+    ) -> "AsyncIterator[_KeyIter[Any]]":
+        for iterable in iterables:
+            iterator = aiter(iterable)
+            try:
+                head = await iterator.__anext__()
+            except StopAsyncIteration:
+                pass
+            else:
+                head_key = await key(head) if key is not None else head
+                yield cls(head, iterator, reverse, head_key, key)
+
+    async def pull_head(self) -> bool:
+        """
+        Pull the next ``head`` element from the iterator and signal success
+        """
+        try:
+            self.head = head = await self.tail.__anext__()
+        except StopAsyncIteration:
+            return False
+        else:
+            self.head_key = await self.key(head) if self.key is not None else head
+            return True
+
+    def __lt__(self, other: "_KeyIter[LT]") -> bool:
+        return self.reverse ^ (self.head_key < other.head_key)
+
+
+@overload
+def merge(
+    *iterables: AnyIterable[LT], key: None = ..., reverse: bool = ...
+) -> AsyncIterator[LT]:
+    pass
+
+
+@overload
+def merge(
+    *iterables: AnyIterable[T],
+    key: Callable[[T], Awaitable[LT]] = ...,
+    reverse: bool = ...
+) -> AsyncIterator[T]:
+    pass
+
+
+@overload
+def merge(
+    *iterables: AnyIterable[T], key: Callable[[T], LT] = ..., reverse: bool = ...
+) -> AsyncIterator[T]:
+    pass
+
+
+async def merge(
+    *iterables: AnyIterable[Any],
+    key: Optional[Callable[[Any], Any]] = None,
+    reverse: bool = False
+) -> AsyncIterator[Any]:
+    """
+    Merge all pre-sorted (async) ``iterables`` into a single sorted iterator
+
+    This works similar to ``sorted(chain(*iterables), key=key, reverse=reverse)`` but
+    operates lazily: at any moment only one item of each iterable is stored for the
+    comparison. This allows merging streams of pre-sorted items, such as timestamped
+    records from multiple sources.
+
+    The optional ``key`` argument specifies a one-argument (async) callable, which
+    provides a substitute for determining the sort order of each item.
+    The special value and default :py:data:`None` represents the identity function,
+    comparing items directly.
+
+    The default sort order is ascending, that is items with ``a < b`` imply ``a``
+    is yielded before ``b``. Use ``reverse=True`` for descending sort order.
+    The ``iterables`` must be pre-sorted in the same order.
+    """
+    a_key = awaitify(key) if key is not None else None
+    # sortable iterators with (reverse) position to ensure stable sort for ties
+    iter_heap: List[Tuple[_KeyIter[Any], int]] = [
+        (itr, idx if not reverse else -idx)
+        async for idx, itr in a_enumerate(
+            _KeyIter.from_iters(iterables, reverse, a_key)
+        )
+    ]
+    try:
+        _heapq.heapify(iter_heap)
+        # there are at least two iterators that need merging
+        while len(iter_heap) > 1:
+            while True:
+                itr, idx = iter_heap[0]
+                yield itr.head
+                if await itr.pull_head():
+                    _heapq.heapreplace(iter_heap, (itr, idx))
+                else:
+                    _heapq.heappop(iter_heap)
+                    break
+        # there is only one iterator left, no need for merging
+        if iter_heap:
+            itr, idx = iter_heap[0]
+            yield itr.head
+            async for item in itr.tail:
+                yield item
+    finally:
+        for itr, _ in iter_heap:
+            if hasattr(itr.tail, "aclose"):
+                await itr.tail.aclose()  # type: ignore
+
+
+class ReverseLT(Generic[LT]):
+    """Helper to reverse ``a < b`` ordering"""
+
+    __slots__ = ("key",)
+
+    def __init__(self, key: LT):
+        self.key = key
+
+    def __lt__(self, other: "ReverseLT[LT]") -> bool:
+        return other.key < self.key
+
+
+# Python's heapq provides a *min*-heap
+# When finding the n largest items, heapq tracks the *minimum* item still large enough.
+# In other words, during search we maintain opposite sort order than what is requested.
+# We turn the min-heap into a max-sort in the end.
+async def _largest(
+    iterable: AsyncIterator[T],
+    n: int,
+    key: Callable[[T], Awaitable[LT]],
+    reverse: bool,
+) -> List[T]:
+    ordered: Callable[[SupportsLT], SupportsLT] = (
+        ReverseLT if reverse else lambda x: x  # type: ignore
+    )
+    async with ScopedIter(iterable) as iterator:
+        # assign an ordering to items to solve ties
+        order_sign = -1 if reverse else 1
+        n_heap = [
+            (ordered(await key(item)), index * order_sign, item)
+            async for index, item in a_zip(range(n), borrow(iterator))
+        ]
+        if not n_heap:
+            return []
+        _heapq.heapify(n_heap)
+        worst_key = n_heap[0][0]
+        next_index = n * order_sign
+        async for item in iterator:
+            item_key = ordered(await key(item))
+            if worst_key < item_key:
+                _heapq.heapreplace(n_heap, (item_key, next_index, item))
+                worst_key = n_heap[0][0]
+                next_index += 1 * order_sign
+        n_heap.sort(reverse=True)
+    return [item for _, _, item in n_heap]
+
+
+async def _identity(x: T) -> T:
+    return x
+
+
+async def nlargest(
+    iterable: AsyncIterator[T],
+    n: int,
+    key: Optional[Callable[[Any], Awaitable[Any]]] = None,
+) -> List[T]:
+    """
+    Return a sorted list of the ``n`` largest elements from the (async) iterable
+
+    The optional ``key`` argument specifies a one-argument (async) callable, which
+    provides a substitute for determining the sort order of each item.
+    The special value and default :py:data:`None` represents the identity functions,
+    comparing items directly.
+
+    The result is equivalent to ``sorted(iterable, key=key, reverse=True)[:n]``,
+    but ``iterable`` is consumed lazily and items are discarded eagerly.
+    """
+    a_key: Callable[[Any], Awaitable[Any]] = (
+        awaitify(key) if key is not None else _identity  # type: ignore
+    )
+    return await _largest(iterable=iterable, n=n, key=a_key, reverse=False)
+
+
+async def nsmallest(
+    iterable: AsyncIterator[T],
+    n: int,
+    key: Optional[Callable[[Any], Awaitable[Any]]] = None,
+) -> List[T]:
+    """
+    Return a sorted list of the ``n`` smallest elements from the (async) iterable
+
+    Provides the reverse functionality to :py:func:`~.nlargest`.
+    """
+    a_key: Callable[[Any], Awaitable[Any]] = (
+        awaitify(key) if key is not None else _identity  # type: ignore
+    )
+    return await _largest(iterable=iterable, n=n, key=a_key, reverse=True)
diff --git a/docs/index.rst b/docs/index.rst
index a4e27c2..972133c 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -31,7 +31,7 @@ The missing ``async`` toolbox
     :synopsis: The async standard library
 
 .. toctree::
-   :maxdepth: 1
+   :maxdepth: 2
    :caption: The Async Toolbox
    :hidden:
 
@@ -39,10 +39,11 @@ The missing ``async`` toolbox
    source/api/functools
    source/api/contextlib
    source/api/itertools
+   source/api/heapq
    source/api/asynctools
 
 .. toctree::
-   :maxdepth: 1
+   :maxdepth: 2
    :caption: Guides and Notes
    :hidden:
 
@@ -50,7 +51,7 @@ The missing ``async`` toolbox
    source/glossary
 
 .. toctree::
-   :maxdepth: 1
+   :maxdepth: 2
    :caption: Development & Maintenance
    :hidden:
 
@@ -93,6 +94,12 @@ with the same name as those of the Python standard library.
     :py:func:`~asyncstdlib.itertools.chain`,
     or :py:func:`~asyncstdlib.itertools.accumulate`.
 
+:py:mod:`asyncstdlib.heapq`
+    Replicates any :py:mod:`heapq` tools that benefit from being asynchronous,
+    which is just :py:func:`~asyncstdlib.heapq.merge`,
+    :py:func:`~asyncstdlib.heapq.nlargest`, and
+    :py:func:`~asyncstdlib.heapq.nsmallest`.
+
 For simplicity, the :py:mod:`asyncstdlib` namespace also exposes all individual
 functions and classes directly.
 For example, :py:mod:`asyncstdlib.builtins.enumerate` is also available
diff --git a/docs/source/api/heapq.rst b/docs/source/api/heapq.rst
new file mode 100644
index 0000000..e3ac192
--- /dev/null
+++ b/docs/source/api/heapq.rst
@@ -0,0 +1,33 @@
+=================
+The heapq library
+=================
+
+.. py:module:: asyncstdlib.heapq
+    :synopsis: async heapq variants
+
+The :py:mod:`asyncstdlib.heapq` library implements
+Python's :py:mod:`heapq` for (async) functions and (async) iterables.
+
+.. versionadded:: 3.10.3
+
+This module does not re-implement the functions to maintain a heap data structure.
+Since Python's :py:mod:`heapq` module does not support an internal ``key`` function
+but relies on ``(key, item)`` pairs as needed,
+the same interface can be used for ``async`` key function.
+For example, an ``async`` key function would be used as
+``heappush(heap, (await key_func(item), item))`` instead of
+``heappush(heap, (key_func(item), item))``.
+
+Iterator merging
+================
+
+.. autofunction:: merge(*iterables: (async) iter T, key: (T) → (await) Any = None, reverse: bool = False)
+    :async-for: :T
+
+Iterator selecting
+==================
+
+.. autofunction:: nlargest(*iterables: (async) iter T, n: int, key: (T) → (await) Any = None) -> [T, ...]
+
+.. autofunction:: nsmallest(*iterables: (async) iter T, n: int, key: (T) → (await) Any = None) -> [T, ...]
+


Test patch:
diff --git a/unittests/test_heapq.py b/unittests/test_heapq.py
new file mode 100644
index 0000000..201b684
--- /dev/null
+++ b/unittests/test_heapq.py
@@ -0,0 +1,88 @@
+import heapq
+
+import pytest
+import random
+
+import asyncstdlib as a
+
+from .utility import sync, asyncify, awaitify
+
+
+MERGE_SAMPLES = [
+    [[1, 2], [3, 4]],
+    [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+    [[1, 4, 7], [2, 5, 8], [3, 6, 9]],
+    [[1], [2, 3], [4, 5, 6], [7, 8, 9]],
+    [[1], [2, 4, 6], [3], [5, 7, 8, 9]],
+    [[1, 2, 3, 4] for _ in range(3)],
+    [sorted(random.random() for _ in range(5)) for _ in range(5)],
+    [[], []],
+    [[]],
+]
+
+
+@pytest.mark.parametrize("samples", MERGE_SAMPLES)
+@pytest.mark.parametrize("reverse", [False, True])
+@sync
+async def test_merge_stdlib(samples, reverse):
+    """Compare `heapq.merge` against stdlib implementation"""
+    samples = samples if not reverse else [sample[::-1] for sample in samples]
+    assert list(heapq.merge(*samples, reverse=reverse)) == [
+        item async for item in a.merge(*map(asyncify, samples), reverse=reverse)
+    ]
+
+
+@pytest.mark.parametrize("samples", MERGE_SAMPLES)
+@pytest.mark.parametrize("reverse", [False, True])
+@sync
+async def test_merge_stdlib_key(samples, reverse):
+    """Compare `heapq.merge` with key against stdlib implementation"""
+    # use a key that reverses the result => must reverse input
+    samples = samples if reverse else [sample[::-1] for sample in samples]
+    assert list(heapq.merge(*samples, key=lambda x: -x, reverse=reverse)) == [
+        item
+        async for item in a.merge(
+            *map(asyncify, samples), key=awaitify(lambda x: -x), reverse=reverse
+        )
+    ]
+
+
+MINMAX_SAMPLES = [
+    [0, 3, 1, 4, 2, 5],
+    [0, -3, 1, -4, -2, 5],
+    # [i * 1.1 for i in range(20)],
+    # [random.randint(-5, 5) in range(20)],
+    # [i * i if i % 3 else -i * i for i in range(2000)],
+]
+
+
+@pytest.mark.parametrize("sample", MINMAX_SAMPLES)
+@pytest.mark.parametrize("n", [0, 1, 2, 10, 100, 400, 999, 1000, 1100])
+@sync
+async def test_nsmallest_stdlib(sample, n):
+    assert heapq.nsmallest(n, sample) == await a.nsmallest(asyncify(sample), n)
+
+
+@pytest.mark.parametrize("sample", MINMAX_SAMPLES)
+@pytest.mark.parametrize("n", [0, 1, 2, 10, 100, 400, 999, 1000, 1100])
+@sync
+async def test_nlargest_stdlib(sample, n):
+    assert heapq.nlargest(n, sample) == await a.nlargest(asyncify(sample), n)
+
+
+@pytest.mark.parametrize("sample", MINMAX_SAMPLES)
+@pytest.mark.parametrize("n", [0, 1, 2, 10, 100, 400, 999, 1000, 1100])
+@sync
+async def test_nsmallest_stdlib_key(sample, n):
+    assert heapq.nsmallest(n, sample, key=lambda x: -x) == await a.nsmallest(
+        asyncify(sample), n, key=lambda x: -x
+    )
+
+
+@pytest.mark.parametrize("sample", MINMAX_SAMPLES)
+@pytest.mark.parametrize("n", [0, 1, 2, 10, 100, 400, 999, 1000, 1100])
+@sync
+async def test_nlargest_stdlib_key(sample, n):
+    assert heapq.nlargest(n, sample, key=lambda x: -x) == await a.nlargest(
+        asyncify(sample), n, key=lambda x: -x
+    )