[Python-checkins] r61369 - in python/trunk: Doc/library/heapq.rst Lib/heapq.py Lib/test/test_heapq.py Misc/NEWS Modules/_heapqmodule.c

[raymond.hettinger]

Author: raymond.hettinger
Date: Thu Mar 13 20:03:51 2008
New Revision: 61369

Modified:
   python/trunk/Doc/library/heapq.rst
   python/trunk/Lib/heapq.py
   python/trunk/Lib/test/test_heapq.py
   python/trunk/Misc/NEWS
   python/trunk/Modules/_heapqmodule.c
Log:
Issue 2274:  Add heapq.heappushpop().

Modified: python/trunk/Doc/library/heapq.rst
==============================================================================
--- python/trunk/Doc/library/heapq.rst	(original)
+++ python/trunk/Doc/library/heapq.rst	Thu Mar 13 20:03:51 2008
@@ -45,6 +45,13 @@
    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)
 

Modified: python/trunk/Lib/heapq.py
==============================================================================
--- python/trunk/Lib/heapq.py	(original)
+++ python/trunk/Lib/heapq.py	Thu Mar 13 20:03:51 2008
@@ -127,7 +127,7 @@
 """
 
 __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 @@
     _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 @@
 
 # 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
 

Modified: python/trunk/Lib/test/test_heapq.py
==============================================================================
--- python/trunk/Lib/test/test_heapq.py	(original)
+++ python/trunk/Lib/test/test_heapq.py	Thu Mar 13 20:03:51 2008
@@ -107,6 +107,34 @@
         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):

Modified: python/trunk/Misc/NEWS
==============================================================================
--- python/trunk/Misc/NEWS	(original)
+++ python/trunk/Misc/NEWS	Thu Mar 13 20:03:51 2008
@@ -491,6 +491,8 @@
 Library
 -------
 
+- #2274 Add heapq.heappushpop().
+
 - Add inspect.isabstract(object) to fix bug #2223
 
 - Add a __format__ method to Decimal, to support PEP 3101.

Modified: python/trunk/Modules/_heapqmodule.c
==============================================================================
--- python/trunk/Modules/_heapqmodule.c	(original)
+++ python/trunk/Modules/_heapqmodule.c	Thu Mar 13 20:03:51 2008
@@ -162,6 +162,11 @@
 {
 	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 @@
             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 @@
 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,