Fix stability of heapq's nlargest() and nsmallest().

diff --git a/Lib/heapq.py b/Lib/heapq.py
index 04725cdef46de4c8e2c019de4c732211c0e01f98..753c3b7ec19891653baf94920130e9d7d10bf794 100644 (file)
--- a/Lib/heapq.py
+++ b/Lib/heapq.py
@@ -130,7 +130,7 @@ __all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'nlargest',
            'nsmallest']
 
 from itertools import islice, repeat, count, imap, izip, tee
-from operator import itemgetter
+from operator import itemgetter, neg
 import bisect
 
 def heappush(heap, item):
@@ -315,8 +315,6 @@ def nsmallest(n, iterable, key=None):
 
     Equivalent to:  sorted(iterable, key=key)[:n]
     """
-    if key is None:
-        return _nsmallest(n, iterable)
     in1, in2 = tee(iterable)
     it = izip(imap(key, in1), count(), in2)                 # decorate
     result = _nsmallest(n, it)
@@ -328,10 +326,8 @@ def nlargest(n, iterable, key=None):
 
     Equivalent to:  sorted(iterable, key=key, reverse=True)[:n]
     """
-    if key is None:
-        return _nlargest(n, iterable)
     in1, in2 = tee(iterable)
-    it = izip(imap(key, in1), count(), in2)                 # decorate
+    it = izip(imap(key, in1), imap(neg, count()), in2)      # decorate
     result = _nlargest(n, it)
     return map(itemgetter(2), result)                       # undecorate
 

diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 2da4f8c49acaa0ef3c14dd6a10936259f3da3ba7..e9f27982bcade28b2b932c3df60d9b64f8625492 100644 (file)
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -104,20 +104,20 @@ class TestHeap(unittest.TestCase):
             self.assertEqual(heap_sorted, sorted(data))
 
     def test_nsmallest(self):
-        data = [random.randrange(2000) for i in range(1000)]
-        f = lambda x:  x * 547 % 2000
-        for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
-            self.assertEqual(nsmallest(n, data), sorted(data)[:n])
-            self.assertEqual(nsmallest(n, data, key=f),
-                             sorted(data, key=f)[:n])
+        data = [(random.randrange(2000), i) for i in range(1000)]
+        for f in (None, lambda x:  x[0] * 547 % 2000):
+            for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
+                self.assertEqual(nsmallest(n, data), sorted(data)[:n])
+                self.assertEqual(nsmallest(n, data, key=f),
+                                 sorted(data, key=f)[:n])
 
     def test_nlargest(self):
-        data = [random.randrange(2000) for i in range(1000)]
-        f = lambda x:  x * 547 % 2000
-        for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
-            self.assertEqual(nlargest(n, data), sorted(data, reverse=True)[:n])
-            self.assertEqual(nlargest(n, data, key=f),
-                             sorted(data, key=f, reverse=True)[:n])
+        data = [(random.randrange(2000), i) for i in range(1000)]
+        for f in (None, lambda x:  x[0] * 547 % 2000):
+            for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
+                self.assertEqual(nlargest(n, data), sorted(data, reverse=True)[:n])
+                self.assertEqual(nlargest(n, data, key=f),
+                                 sorted(data, key=f, reverse=True)[:n])
 
 
 #==============================================================================

diff --git a/Misc/NEWS b/Misc/NEWS
index 7ab30ed2685d4f589c6566be661d1efbadec8292..aab3e4372cae6d37347689544ec0591d8288fbe4 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -125,6 +125,8 @@ Extension Modules
 Library
 -------
 
+- Fix sort stability in heapq.nlargest() and nsmallest().
+
 - Patch #1504073: Fix tarfile.open() for mode "r" with a fileobj argument.
 
 - Patch #1262036: Prevent TarFiles from being added to themselves under