Preserve input order for equal-key top-N results

asyncstdlib/heapq.py

@@ -163,6 +163,9 @@ def __init__(self, key: LT):
     def __lt__(self, other: ReverseLT[LT]) -> bool:
         return bool(other.key < self.key)
 
+    def __eq__(self, other: ReverseLT[LT]) -> bool:  # type: ignore[override]
+        return self.key is other.key or self.key == other.key
+
 
 # Python's heapq provides a *min*-heap
 # When finding the n largest items, heapq tracks the *minimum* item still large enough.
@@ -177,7 +180,7 @@ async def _largest(
     ordered: Callable[[LT], LT] = 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
+        order_sign = -1
         n_heap = [
             (ordered(await key(item)), index * order_sign, item)
             async for index, item in a_zip(range(n), borrow(iterator))

unittests/test_heapq.py

@@ -85,3 +85,31 @@ 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
     )
+
+
+@pytest.mark.parametrize("name", ["nlargest", "nsmallest"])
+@pytest.mark.parametrize("keys", [[1, 1, 1], [2, 2, 1, 1, 3, 3], [3, 1, 2, 1, 2, 3]])
+@pytest.mark.parametrize("n", [1, 2, 3, 5, 10])
+@pytest.mark.parametrize("async_key", [False, True])
+@sync
+async def test_nlargest_nsmallest_stable_ties(name, keys, n, async_key):
+    # Dict payloads cannot be ordered: ties must use input position, not the item.
+    values = [{"key": key, "position": index} for index, key in enumerate(keys)]
+
+    def key(value):
+        return value["key"]
+
+    expected = getattr(heapq, name)(n, values, key=key)
+    actual = await getattr(a, name)(
+        asyncify(values), n, key=awaitify(key) if async_key else key
+    )
+    assert actual == expected
+
+
+@pytest.mark.parametrize("name", ["nlargest", "nsmallest"])
+@sync
+async def test_nlargest_nsmallest_identical_key(name):
+    key = float("nan")
+    values = [{"position": index} for index in range(3)]
+    expected = getattr(heapq, name)(2, values, key=lambda value: key)
+    assert await getattr(a, name)(values, 2, key=lambda value: key) == expected