From 8aaddc87669c38f88ee88daa2b244b585d8d4394 Mon Sep 17 00:00:00 2001 From: Facundo Batista Date: Tue, 7 Jul 2026 16:05:21 -0300 Subject: [PATCH 1/4] Implemented the test cases and Pyhon classes for the basic functionality. --- Lib/heapq.py | 124 ++++++++++++++++++++++++++++- Lib/test/test_heapq.py | 175 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 297 insertions(+), 2 deletions(-) diff --git a/Lib/heapq.py b/Lib/heapq.py index a3af6dc05bff377..7e59cbdc8ef8543 100644 --- a/Lib/heapq.py +++ b/Lib/heapq.py @@ -128,7 +128,8 @@ __all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'heappushpop', 'heappush_max', 'heappop_max', 'heapify_max', 'heapreplace_max', - 'heappushpop_max', 'nlargest', 'nsmallest', 'merge'] + 'heappushpop_max', 'nlargest', 'nsmallest', 'merge', + 'MinHeap', 'MaxHeap'] def heappush(heap, item): """Push item onto heap, maintaining the heap invariant.""" @@ -592,6 +593,127 @@ def nlargest(n, iterable, key=None): result.sort(reverse=True) return [elem for (k, order, elem) in result] + +class MinHeap: + """Min-heap: a priority queue that serves its smallest item first. + + Only the ``<`` operator is used to compare items. + + If *iterable* is given, the heap is initialized from its items in linear + time; otherwise it starts empty. + + >>> h = MinHeap([3, 1, 2]) + >>> h.push(0) + >>> [h.pop() for _ in range(len(h))] + [0, 1, 2, 3] + """ + + def __init__(self, iterable=None): + if iterable is None: + self._queue = [] + else: + self._queue = list(iterable) + heapify(self._queue) + + def push(self, item): + """Push item onto the heap, maintaining the heap invariant.""" + heappush(self._queue, item) + + def pop(self): + """Pop and return the smallest item, maintaining the heap invariant. + + Raises IndexError if the heap is empty. + """ + return heappop(self._queue) + + def pushpop(self, item): + """Push item on the heap, then pop and return the smallest item. + + The combined action runs more efficiently than a push() followed by + a separate call to pop(). + """ + return heappushpop(self._queue, item) + + def replace(self, item): + """Pop and return the smallest item, and also push the new item. + + The heap size doesn't change. Raises IndexError if the heap is + empty. This is more efficient than a pop() followed by a push(), + and can be more appropriate when using a fixed-size heap. Note that + the value returned may be larger than item. + """ + return heapreplace(self._queue, item) + + def __len__(self): + return len(self._queue) + + def __bool__(self): + return bool(self._queue) + + def __repr__(self): + return f'{type(self).__name__}({self._queue!r})' + + +class MaxHeap: + """Max-heap: a priority queue that serves its largest item first. + + Only the ``<`` operator is used to compare items. + + If *iterable* is given, the heap is initialized from its items in linear + time; otherwise it starts empty. + + >>> h = MaxHeap([1, 3, 2]) + >>> h.push(4) + >>> [h.pop() for _ in range(len(h))] + [4, 3, 2, 1] + """ + + def __init__(self, iterable=None): + if iterable is None: + self._queue = [] + else: + self._queue = list(iterable) + heapify_max(self._queue) + + def push(self, item): + """Push item onto the heap, maintaining the heap invariant.""" + heappush_max(self._queue, item) + + def pop(self): + """Pop and return the largest item, maintaining the heap invariant. + + Raises IndexError if the heap is empty. + """ + return heappop_max(self._queue) + + def pushpop(self, item): + """Push item on the heap, then pop and return the largest item. + + The combined action runs more efficiently than a push() followed by + a separate call to pop(). + """ + return heappushpop_max(self._queue, item) + + def replace(self, item): + """Pop and return the largest item, and also push the new item. + + The heap size doesn't change. Raises IndexError if the heap is + empty. This is more efficient than a pop() followed by a push(), + and can be more appropriate when using a fixed-size heap. Note that + the value returned may be smaller than item. + """ + return heapreplace_max(self._queue, item) + + def __len__(self): + return len(self._queue) + + def __bool__(self): + return bool(self._queue) + + def __repr__(self): + return f'{type(self).__name__}({self._queue!r})' + + # If available, use C implementation try: from _heapq import * diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index d6623fee9bb2b43..d1008a619833277 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -5,7 +5,7 @@ import doctest from test.support import import_helper -from unittest import TestCase, skipUnless +from unittest import TestCase, mock, skipUnless from operator import itemgetter py_heapq = import_helper.import_fresh_module('heapq', blocked=['_heapq']) @@ -637,5 +637,178 @@ class TestErrorHandlingC(TestErrorHandling, TestCase): module = c_heapq +#============================================================================== +# Object-oriented interface: MinHeap and MaxHeap. +# +# The classes add no new heap algorithm; they are thin wrappers that delegate +# to the module-level functions. They are tested in two ways: +# +# * Functional "sanity check" tests, run against both the pure-Python and the +# C classes (like the rest of this file). They are end-to-end, so a +# mis-wired method (say, a MaxHeap that pushed with the min-heap routine) +# changes the observable output and is caught even for the C class. +# +# * Delegation tests that patch the module-level functions and check that each +# method forwards to the right one. These only apply to the pure-Python +# class, since the C class calls the underlying routines directly and cannot +# be intercepted; its wiring is covered by the functional tests instead. + + +class HeapClassSanityTest: + """Functional sanity checks shared by the MinHeap/MaxHeap test cases. + + Subclasses set: + + * *cls* -- the heap class under test (e.g. ``py_heapq.MinHeap``) + * *ordered* -- a small list already satisfying the heap invariant + * *reverse* -- ``True`` for max-heaps (the order ``pop()`` yields items in) + """ + + def in_heap_order(self, data): + return sorted(data, reverse=self.reverse) + + def test_construct_empty(self): + h = self.cls() + self.assertEqual(len(h), 0) + self.assertFalse(h) + + def test_construct_from_iterable(self): + data = [random.randrange(1000) for _ in range(100)] + # Accept any iterable, not just a sequence. + h = self.cls(iter(data)) + self.assertEqual(len(h), len(data)) + drained = [h.pop() for _ in range(len(data))] + self.assertEqual(drained, self.in_heap_order(data)) + + def test_push_then_pop_sorts(self): + data = [random.randrange(1000) for _ in range(100)] + h = self.cls() + for item in data: + self.assertIsNone(h.push(item)) + self.assertEqual(len(h), len(data)) + drained = [h.pop() for _ in range(len(data))] + self.assertEqual(drained, self.in_heap_order(data)) + + def test_len_and_bool(self): + h = self.cls() + self.assertEqual(len(h), 0) + self.assertFalse(h) + h.push(1) + self.assertEqual(len(h), 1) + self.assertTrue(h) + + def test_pop_empty_raises(self): + h = self.cls() + self.assertRaises(IndexError, h.pop) + + def test_pushpop(self): + h = self.cls(self.ordered) + result = h.pushpop(2) + expected = self.in_heap_order(list(self.ordered) + [2])[0] + self.assertEqual(result, expected) + self.assertEqual(len(h), len(self.ordered)) + + def test_replace(self): + h = self.cls(self.ordered) + expected = self.in_heap_order(self.ordered)[0] + new_item = -1000 if self.reverse else 1000 + self.assertEqual(h.replace(new_item), expected) + self.assertEqual(len(h), len(self.ordered)) + + def test_repr(self): + h = self.cls(self.ordered) + self.assertEqual(repr(h), f'{self.cls.__name__}({self.ordered!r})') + + +class TestMinHeapPython(HeapClassSanityTest, TestCase): + cls = py_heapq.MinHeap + ordered = [1, 2, 3] + reverse = False + + +@skipUnless(c_heapq, 'requires _heapq') +class TestMinHeapC(HeapClassSanityTest, TestCase): + cls = c_heapq.MinHeap if c_heapq else None + ordered = [1, 2, 3] + reverse = False + + +class TestMaxHeapPython(HeapClassSanityTest, TestCase): + cls = py_heapq.MaxHeap + ordered = [3, 2, 1] + reverse = True + + +@skipUnless(c_heapq, 'requires _heapq') +class TestMaxHeapC(HeapClassSanityTest, TestCase): + cls = c_heapq.MaxHeap if c_heapq else None + ordered = [3, 2, 1] + reverse = True + + +class HeapClassDelegationTest: + """Check that each method forwards to the right module-level function. + + Delegation only makes sense for the pure-Python class, so these tests + always patch ``py_heapq``. + + Subclasses set *cls* and the names of the functions each method is + expected to forward to (*heapify_func*, *push_func*, *pop_func*, + *pushpop_func* and *replace_func*). + """ + + def test_init_delegates_to_heapify(self): + with mock.patch.object(py_heapq, self.heapify_func) as m: + h = self.cls([3, 1, 2]) + m.assert_called_once_with(h._queue) + + def test_init_empty_does_not_heapify(self): + with mock.patch.object(py_heapq, self.heapify_func) as m: + self.cls() + m.assert_not_called() + + def test_push_delegates(self): + h = self.cls() + with mock.patch.object(py_heapq, self.push_func) as m: + h.push(42) + m.assert_called_once_with(h._queue, 42) + + def test_pop_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, self.pop_func) as m: + h.pop() + m.assert_called_once_with(h._queue) + + def test_pushpop_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, self.pushpop_func) as m: + h.pushpop(42) + m.assert_called_once_with(h._queue, 42) + + def test_replace_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, self.replace_func) as m: + h.replace(42) + m.assert_called_once_with(h._queue, 42) + + +class TestMinHeapDelegation(HeapClassDelegationTest, TestCase): + cls = py_heapq.MinHeap + heapify_func = 'heapify' + push_func = 'heappush' + pop_func = 'heappop' + pushpop_func = 'heappushpop' + replace_func = 'heapreplace' + + +class TestMaxHeapDelegation(HeapClassDelegationTest, TestCase): + cls = py_heapq.MaxHeap + heapify_func = 'heapify_max' + push_func = 'heappush_max' + pop_func = 'heappop_max' + pushpop_func = 'heappushpop_max' + replace_func = 'heapreplace_max' + + if __name__ == "__main__": unittest.main() From 31e1e3d232a80254e37b6d4b1af722bdac54b97b Mon Sep 17 00:00:00 2001 From: Facundo Batista Date: Wed, 8 Jul 2026 18:07:34 -0300 Subject: [PATCH 2/4] Test cases and Python methos for nlargest/nsmallest --- Lib/heapq.py | 32 ++++++++++++++++++++++++++++++++ Lib/test/test_heapq.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/Lib/heapq.py b/Lib/heapq.py index 7e59cbdc8ef8543..5cf7598abb117fd 100644 --- a/Lib/heapq.py +++ b/Lib/heapq.py @@ -644,6 +644,22 @@ def replace(self, item): """ return heapreplace(self._queue, item) + def nsmallest(self, n, key=None): + """Return a list of the n smallest items, smallest first. + + The heap is left unchanged. If *key* is given, it is applied to + each item to determine the ordering. + """ + return nsmallest(n, self._queue, key=key) + + def nlargest(self, n, key=None): + """Return a list of the n largest items, largest first. + + The heap is left unchanged. If *key* is given, it is applied to + each item to determine the ordering. + """ + return nlargest(n, self._queue, key=key) + def __len__(self): return len(self._queue) @@ -704,6 +720,22 @@ def replace(self, item): """ return heapreplace_max(self._queue, item) + def nsmallest(self, n, key=None): + """Return a list of the n smallest items, smallest first. + + The heap is left unchanged. If *key* is given, it is applied to + each item to determine the ordering. + """ + return nsmallest(n, self._queue, key=key) + + def nlargest(self, n, key=None): + """Return a list of the n largest items, largest first. + + The heap is left unchanged. If *key* is given, it is applied to + each item to determine the ordering. + """ + return nlargest(n, self._queue, key=key) + def __len__(self): return len(self._queue) diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index d1008a619833277..2fa4296058b5a1e 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -715,6 +715,26 @@ def test_replace(self): self.assertEqual(h.replace(new_item), expected) self.assertEqual(len(h), len(self.ordered)) + def test_nsmallest(self): + data = [random.randrange(1000) for _ in range(100)] + h = self.cls(data) + self.assertEqual(h.nsmallest(10), sorted(data)[:10]) + # the heap is not modified + self.assertEqual(len(h), len(data)) + + def test_nlargest(self): + data = [random.randrange(1000) for _ in range(100)] + h = self.cls(data) + self.assertEqual(h.nlargest(10), sorted(data, reverse=True)[:10]) + # the heap is not modified + self.assertEqual(len(h), len(data)) + + def test_nsmallest_nlargest_with_key(self): + data = ['a', 'ee', 'ddd', 'bbbb', 'ccccc'] + h = self.cls(data) + self.assertEqual(h.nsmallest(2, key=len), ['a', 'ee']) + self.assertEqual(h.nlargest(2, key=len), ['ccccc', 'bbbb']) + def test_repr(self): h = self.cls(self.ordered) self.assertEqual(repr(h), f'{self.cls.__name__}({self.ordered!r})') @@ -791,6 +811,18 @@ def test_replace_delegates(self): h.replace(42) m.assert_called_once_with(h._queue, 42) + def test_nsmallest_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, 'nsmallest') as m: + h.nsmallest(2, key=str) + m.assert_called_once_with(2, h._queue, key=str) + + def test_nlargest_delegates(self): + h = self.cls([1, 2, 3]) + with mock.patch.object(py_heapq, 'nlargest') as m: + h.nlargest(2, key=str) + m.assert_called_once_with(2, h._queue, key=str) + class TestMinHeapDelegation(HeapClassDelegationTest, TestCase): cls = py_heapq.MinHeap From b8b1667815ce8ea071b38a3dabe836ba702e8d19 Mon Sep 17 00:00:00 2001 From: Facundo Batista Date: Wed, 8 Jul 2026 18:45:48 -0300 Subject: [PATCH 3/4] Improved tests --- Lib/test/test_heapq.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index 2fa4296058b5a1e..01e575be279a773 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -17,6 +17,8 @@ # Add max-heap variants func_names += [func + '_max' for func in func_names] +class_names = ['MinHeap', 'MaxHeap'] + class TestModules(TestCase): def test_py_functions(self): for fname in func_names: @@ -27,6 +29,15 @@ def test_c_functions(self): for fname in func_names: self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq', fname) + def test_py_classes(self): + # MinHeap and MaxHeap are only implemented in pure Python; there is no + # C accelerator for them. They rely on the module-level functions, + # which are the C versions when _heapq is available. + for cname in class_names: + self.assertEqual(getattr(py_heapq, cname).__module__, 'heapq') + if c_heapq is not None: + self.assertEqual(getattr(c_heapq, cname).__module__, 'heapq') + def load_tests(loader, tests, ignore): # The 'merge' function has examples in its docstring which we should test @@ -697,6 +708,9 @@ def test_len_and_bool(self): self.assertEqual(len(h), 1) self.assertTrue(h) + def test_construct_non_iterable_raises(self): + self.assertRaises(TypeError, self.cls, 10) + def test_pop_empty_raises(self): h = self.cls() self.assertRaises(IndexError, h.pop) From 4d1581ad44ff1c2efb4ff4397c5be45fc4bf1351 Mon Sep 17 00:00:00 2001 From: Facundo Batista Date: Wed, 15 Jul 2026 15:46:22 -0300 Subject: [PATCH 4/4] Docs and what's new --- Doc/library/heapq.rst | 86 +++++++++++++++++++++++++++++++++++++++++++ Doc/whatsnew/3.16.rst | 7 ++++ 2 files changed, 93 insertions(+) diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst index 26cffa7c643028a..3349fc6c8f2a0d6 100644 --- a/Doc/library/heapq.rst +++ b/Doc/library/heapq.rst @@ -50,6 +50,83 @@ To create a heap, use a list initialized as ``[]``, or transform an existing lis into a min-heap or max-heap using the :func:`heapify` or :func:`heapify_max` functions, respectively. +Alternatively, the :class:`MinHeap` and :class:`MaxHeap` classes offer an +object-oriented interface, described next; the module-level functions that they +build upon are documented afterwards. + + +Heap classes +------------ + +The :class:`MinHeap` and :class:`MaxHeap` classes wrap the module's heap +operations in an object that owns its own storage. Compared with the functions +below, there is no separate list to pass to every call or to accidentally +break the heap invariant with an unrelated list operation. + +.. class:: MinHeap(iterable=None) + MaxHeap(iterable=None) + + Create a new heap. If *iterable* is given, the heap is initialized from its + items in linear time; otherwise it starts empty. + + A :class:`MinHeap` keeps its smallest item at the top, while a + :class:`MaxHeap` keeps its largest item at the top. The two classes share + the same interface and differ only in that ordering --- that is, in which + item :meth:`~MinHeap.pop` removes and which item the other methods treat as + the top. As with the functions below, only the ``<`` operator is used to + compare items. + + The methods documented below belong to both classes; a :class:`MinHeap` is + used in the descriptions, with *MaxHeap* behaving the same way once + "smallest" is read as "largest". + + .. method:: push(item) + + Push the value *item* onto the heap, maintaining the heap invariant. + + .. method:: pop() + + Pop and return the smallest item from the heap, maintaining the heap + invariant. If the heap is empty, :exc:`IndexError` is raised. + + .. method:: pushpop(item) + + Push *item* on the heap, then pop and return the smallest item. The + combined action runs more efficiently than :meth:`push` followed by a + separate call to :meth:`pop`. + + .. method:: replace(item) + + Pop and return the smallest item from the heap, and also push the new + *item*. The heap size doesn't change. If the heap is empty, + :exc:`IndexError` is raised. + + This one step operation is more efficient than a :meth:`pop` followed by + :meth:`push` and can be more appropriate when using a fixed-size heap. + The value returned may be larger than the *item* added; consider using + :meth:`pushpop` instead if that is not desired. + + .. method:: nsmallest(n, key=None) + + Return a list with the *n* smallest items from the heap, from smallest to + largest, without modifying the heap. Behaves like the module-level + :func:`nsmallest` function applied to the heap's items. + + .. method:: nlargest(n, key=None) + + Return a list with the *n* largest items from the heap, from largest to + smallest, without modifying the heap. Behaves like the module-level + :func:`nlargest` function applied to the heap's items. + + In addition, ``len(heap)`` returns the number of items on the heap, and an + empty heap is falsy, so ``bool(heap)`` is ``False`` when it holds no items. + + .. versionadded:: 3.16 + + +Heap functions +-------------- + The following functions are provided for min-heaps: @@ -214,6 +291,15 @@ time:: This is similar to ``sorted(iterable)``, but unlike :func:`sorted`, this implementation is not stable. +Using the object-oriented interface, the same sort becomes:: + + >>> def heapsort(iterable): + ... heap = MinHeap(iterable) + ... return [heap.pop() for _ in range(len(heap))] + ... + >>> heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0]) + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + Heap elements can be tuples. This is useful for assigning comparison values (such as task priorities) alongside the main record being tracked:: diff --git a/Doc/whatsnew/3.16.rst b/Doc/whatsnew/3.16.rst index 1a73a79a58b78b1..19ca5d9db1710f3 100644 --- a/Doc/whatsnew/3.16.rst +++ b/Doc/whatsnew/3.16.rst @@ -221,6 +221,13 @@ gzip which is passed on to the constructor of the :class:`~gzip.GzipFile` class. (Contributed by Marin Misur in :gh:`91372`.) +heapq +----- + +* Add the :class:`~heapq.MinHeap` and :class:`~heapq.MaxHeap` classes, an + object-oriented interface to the module's heap functions. + (Contributed by Facundo Batista in :gh:`137200`.) + io --