Issue 2274: Add heapq.heappushpop().

5 files changed, 95 insertions, 2 deletions

diff --git a/Doc/library/heapq.rst b/Doc/library/heapq.rst
index 115d223..1168fb6 100644
--- 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 39e3800..23f5fcb 100644
--- 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 7f6e4b5..fec027e 100644
--- 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 bd3a6cc..4d0b9f0 100644
--- 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 f970cae..703742e 100644
--- 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;
 
@@ -196,6 +201,48 @@ this routine unless written as part of a conditional replacement:\n\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)
 {
 	Py_ssize_t i, n;
@@ -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,