Issue 2274: Add heapq.heappushpop().

diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst
index 115d22369217b0b6e74770a199582f10c7be409b..1168fb688a2ce2b0ad32b3e1f359f914425dfc92 100644 (file)
--- a/Doc/library/heapq.rst
+++ b/Doc/library/heapq.rst
@@ -45,6 +45,13 @@ The following functions are provided:
    Pop and return the smallest item from the *heap*, maintaining the heap
    invariant.  If the heap is empty, :exc:`IndexError` is raised.
 
+.. function:: heappushpop(heap, item)
+
+   Push *item* on the heap, then pop and return the smallest item from the
+   *heap*.  The combined action runs more efficiently than :func:`heappush`
+   followed by a separate call to :func:`heappop`.
+
+   .. versionadded:: 2.6
 
 .. function:: heapify(x)
 

diff --git a/Lib/heapq.py b/Lib/heapq.py
index 39e38008822a0add0e47fa0e989d9316458cc0cf..23f5fcbfda078ac80442b22099998520e330cdc2 100644 (file)
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -127,7 +127,7 @@ From all times, sorting has always been a Great Art! :-)
 """
 
 __all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge',
-           'nlargest', 'nsmallest']
+           'nlargest', 'nsmallest', 'heappushpop']
 
 from itertools import islice, repeat, count, imap, izip, tee
 from operator import itemgetter, neg
@@ -165,6 +165,13 @@ def heapreplace(heap, item):
     _siftup(heap, 0)
     return returnitem
 
+def heappushpop(heap, item):
+    """Fast version of a heappush followed by a heappop."""
+    if heap and item > heap[0]:
+        item, heap[0] = heap[0], item
+        _siftup(heap, 0)
+    return item
+
 def heapify(x):
     """Transform list into a heap, in-place, in O(len(heap)) time."""
     n = len(x)
@@ -304,7 +311,7 @@ def _siftup(heap, pos):
 
 # If available, use C implementation
 try:
-    from _heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest
+    from _heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest, heappushpop
 except ImportError:
     pass
 

diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 7f6e4b59e707af1228601b3e339d9f1fb8a13004..fec027e3b659fbc2c97dbc9256c7c3256e969974 100644 (file)
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -107,6 +107,34 @@ class TestHeap(unittest.TestCase):
         self.assertRaises(TypeError, self.module.heapreplace, None, None)
         self.assertRaises(IndexError, self.module.heapreplace, [], None)
 
+    def test_nbest_with_pushpop(self):
+        data = [random.randrange(2000) for i in range(1000)]
+        heap = data[:10]
+        self.module.heapify(heap)
+        for item in data[10:]:
+            self.module.heappushpop(heap, item)
+        self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:])
+        self.assertEqual(self.module.heappushpop([], 'x'), 'x')
+
+    def test_heappushpop(self):
+        h = []
+        x = self.module.heappushpop(h, 10)
+        self.assertEqual((h, x), ([], 10))
+
+        h = [10]
+        x = self.module.heappushpop(h, 10.0)
+        self.assertEqual((h, x), ([10], 10.0))
+        self.assertEqual(type(h[0]), int)
+        self.assertEqual(type(x), float)
+
+        h = [10];
+        x = self.module.heappushpop(h, 9)
+        self.assertEqual((h, x), ([10], 9))
+
+        h = [10];
+        x = self.module.heappushpop(h, 11)
+        self.assertEqual((h, x), ([11], 10))
+
     def test_heapsort(self):
         # Exercise everything with repeated heapsort checks
         for trial in xrange(100):

diff --git a/Misc/NEWS b/Misc/NEWS
index bd3a6cca0dc682740a8c7329d2bb777bac49356a..4d0b9f0fb5215c29a3b3cbffe2fa5a52ab07dfdf 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -491,6 +491,8 @@ Core and builtins
 Library
 -------
 
+- #2274 Add heapq.heappushpop().
+
 - Add inspect.isabstract(object) to fix bug #2223
 
 - Add a __format__ method to Decimal, to support PEP 3101.

diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index f970cae91377cc683aa82f3277d900e236ff51e1..703742e26085a730519f6ed31e0edf6c76b14ad6 100644 (file)
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -162,6 +162,11 @@ heapreplace(PyObject *self, PyObject *args)
 {
        PyObject *heap, *item, *returnitem;
 
+       if (Py_Py3kWarningFlag &&
+           PyErr_Warn(PyExc_DeprecationWarning, 
+                      "In 3.x, heapreplace() was removed. Use heappushpop() instead.") < 0)
+               return NULL;
+
        if (!PyArg_UnpackTuple(args, "heapreplace", 2, 2, &heap, &item))
                return NULL;
 
@@ -195,6 +200,48 @@ this routine unless written as part of a conditional replacement:\n\n\
         if item > heap[0]:\n\
             item = heapreplace(heap, item)\n");
 
+static PyObject *
+heappushpop(PyObject *self, PyObject *args)
+{
+       PyObject *heap, *item, *returnitem;
+       int cmp;
+
+       if (!PyArg_UnpackTuple(args, "heappushpop", 2, 2, &heap, &item))
+               return NULL;
+
+       if (!PyList_Check(heap)) {
+               PyErr_SetString(PyExc_TypeError, "heap argument must be a list");
+               return NULL;
+       }
+
+       if (PyList_GET_SIZE(heap) < 1) {
+               Py_INCREF(item);
+               return item;
+       }
+
+       cmp = PyObject_RichCompareBool(item, PyList_GET_ITEM(heap, 0), Py_LE);
+       if (cmp == -1)
+               return NULL;
+       if (cmp == 1) {
+               Py_INCREF(item);
+               return item;
+       }
+
+       returnitem = PyList_GET_ITEM(heap, 0);
+       Py_INCREF(item);
+       PyList_SET_ITEM(heap, 0, item);
+       if (_siftup((PyListObject *)heap, 0) == -1) {
+               Py_DECREF(returnitem);
+               return NULL;
+       }
+       return returnitem;
+}
+
+PyDoc_STRVAR(heappushpop_doc,
+"Push item on the heap, then pop and return the smallest item\n\
+from the heap. The combined action runs more efficiently than\n\
+heappush() followed by a separate call to heappop().");
+
 static PyObject *
 heapify(PyObject *self, PyObject *heap)
 {
@@ -468,6 +515,8 @@ Equivalent to:  sorted(iterable)[:n]\n");
 static PyMethodDef heapq_methods[] = {
        {"heappush",    (PyCFunction)heappush,          
                METH_VARARGS,   heappush_doc},
+       {"heappushpop", (PyCFunction)heappushpop,               
+               METH_VARARGS,   heappushpop_doc},
        {"heappop",     (PyCFunction)heappop,
                METH_O,         heappop_doc},
        {"heapreplace", (PyCFunction)heapreplace,