Lib/test/test_heapq.py

   1 """Unittests for heapq."""
   2
   3 import random
   4 import unittest
   5 from test import test_support
   6 import sys
   7
   8 # We do a bit of trickery here to be able to test both the C implementation
   9 # and the Python implementation of the module.
  10 import heapq as c_heapq
  11 py_heapq = test_support.import_fresh_module('heapq', blocked=['_heapq'])
  12
  13 class TestHeap(unittest.TestCase):
  14     module = None
  15
  16     def test_push_pop(self):
  17         # 1) Push 256 random numbers and pop them off, verifying all's OK.
  18         heap = []
  19         data = []
  20         self.check_invariant(heap)
  21         for i in range(256):
  22             item = random.random()
  23             data.append(item)
  24             self.module.heappush(heap, item)
  25             self.check_invariant(heap)
  26         results = []
  27         while heap:
  28             item = self.module.heappop(heap)
  29             self.check_invariant(heap)
  30             results.append(item)
  31         data_sorted = data[:]
  32         data_sorted.sort()
  33         self.assertEqual(data_sorted, results)
  34         # 2) Check that the invariant holds for a sorted array
  35         self.check_invariant(results)
  36
  37         self.assertRaises(TypeError, self.module.heappush, [])
  38         try:
  39             self.assertRaises(TypeError, self.module.heappush, None, None)
  40             self.assertRaises(TypeError, self.module.heappop, None)
  41         except AttributeError:
  42             pass
  43
  44     def check_invariant(self, heap):
  45         # Check the heap invariant.
  46         for pos, item in enumerate(heap):
  47             if pos: # pos 0 has no parent
  48                 parentpos = (pos-1) >> 1
  49                 self.assertTrue(heap[parentpos] <= item)
  50
  51     def test_heapify(self):
  52         for size in range(30):
  53             heap = [random.random() for dummy in range(size)]
  54             self.module.heapify(heap)
  55             self.check_invariant(heap)
  56
  57         self.assertRaises(TypeError, self.module.heapify, None)
  58
  59     def test_naive_nbest(self):
  60         data = [random.randrange(2000) for i in range(1000)]
  61         heap = []
  62         for item in data:
  63             self.module.heappush(heap, item)
  64             if len(heap) > 10:
  65                 self.module.heappop(heap)
  66         heap.sort()
  67         self.assertEqual(heap, sorted(data)[-10:])
  68
  69     def heapiter(self, heap):
  70         # An iterator returning a heap's elements, smallest-first.
  71         try:
  72             while 1:
  73                 yield self.module.heappop(heap)
  74         except IndexError:
  75             pass
  76
  77     def test_nbest(self):
  78         # Less-naive "N-best" algorithm, much faster (if len(data) is big
  79         # enough <wink>) than sorting all of data.  However, if we had a max
  80         # heap instead of a min heap, it could go faster still via
  81         # heapify'ing all of data (linear time), then doing 10 heappops
  82         # (10 log-time steps).
  83         data = [random.randrange(2000) for i in range(1000)]
  84         heap = data[:10]
  85         self.module.heapify(heap)
  86         for item in data[10:]:
  87             if item > heap[0]:  # this gets rarer the longer we run
  88                 self.module.heapreplace(heap, item)
  89         self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:])
  90
  91         self.assertRaises(TypeError, self.module.heapreplace, None)
  92         self.assertRaises(TypeError, self.module.heapreplace, None, None)
  93         self.assertRaises(IndexError, self.module.heapreplace, [], None)
  94
  95     def test_nbest_with_pushpop(self):
  96         data = [random.randrange(2000) for i in range(1000)]
  97         heap = data[:10]
  98         self.module.heapify(heap)
  99         for item in data[10:]:
 100             self.module.heappushpop(heap, item)
 101         self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:])
 102         self.assertEqual(self.module.heappushpop([], 'x'), 'x')
 103
 104     def test_heappushpop(self):
 105         h = []
 106         x = self.module.heappushpop(h, 10)
 107         self.assertEqual((h, x), ([], 10))
 108
 109         h = [10]
 110         x = self.module.heappushpop(h, 10.0)
 111         self.assertEqual((h, x), ([10], 10.0))
 112         self.assertEqual(type(h[0]), int)
 113         self.assertEqual(type(x), float)
 114
 115         h = [10];
 116         x = self.module.heappushpop(h, 9)
 117         self.assertEqual((h, x), ([10], 9))
 118
 119         h = [10];
 120         x = self.module.heappushpop(h, 11)
 121         self.assertEqual((h, x), ([11], 10))
 122
 123     def test_heapsort(self):
 124         # Exercise everything with repeated heapsort checks
 125         for trial in xrange(100):
 126             size = random.randrange(50)
 127             data = [random.randrange(25) for i in range(size)]
 128             if trial & 1:     # Half of the time, use heapify
 129                 heap = data[:]
 130                 self.module.heapify(heap)
 131             else:             # The rest of the time, use heappush
 132                 heap = []
 133                 for item in data:
 134                     self.module.heappush(heap, item)
 135             heap_sorted = [self.module.heappop(heap) for i in range(size)]
 136             self.assertEqual(heap_sorted, sorted(data))
 137
 138     def test_merge(self):
 139         inputs = []
 140         for i in xrange(random.randrange(5)):
 141             row = sorted(random.randrange(1000) for j in range(random.randrange(10)))
 142             inputs.append(row)
 143         self.assertEqual(sorted(chain(*inputs)), list(self.module.merge(*inputs)))
 144         self.assertEqual(list(self.module.merge()), [])
 145
 146     def test_merge_stability(self):
 147         class Int(int):
 148             pass
 149         inputs = [[], [], [], []]
 150         for i in range(20000):
 151             stream = random.randrange(4)
 152             x = random.randrange(500)
 153             obj = Int(x)
 154             obj.pair = (x, stream)
 155             inputs[stream].append(obj)
 156         for stream in inputs:
 157             stream.sort()
 158         result = [i.pair for i in self.module.merge(*inputs)]
 159         self.assertEqual(result, sorted(result))
 160
 161     def test_nsmallest(self):
 162         data = [(random.randrange(2000), i) for i in range(1000)]
 163         for f in (None, lambda x:  x[0] * 547 % 2000):
 164             for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
 165                 self.assertEqual(self.module.nsmallest(n, data), sorted(data)[:n])
 166                 self.assertEqual(self.module.nsmallest(n, data, key=f),
 167                                  sorted(data, key=f)[:n])
 168
 169     def test_nlargest(self):
 170         data = [(random.randrange(2000), i) for i in range(1000)]
 171         for f in (None, lambda x:  x[0] * 547 % 2000):
 172             for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
 173                 self.assertEqual(self.module.nlargest(n, data),
 174                                  sorted(data, reverse=True)[:n])
 175                 self.assertEqual(self.module.nlargest(n, data, key=f),
 176                                  sorted(data, key=f, reverse=True)[:n])
 177
 178 class TestHeapPython(TestHeap):
 179     module = py_heapq
 180
 181     # As an early adopter, we sanity check the
 182     # test_support.import_fresh_module utility function
 183     def test_pure_python(self):
 184         self.assertFalse(sys.modules['heapq'] is self.module)
 185         self.assertTrue(hasattr(self.module.heapify, 'func_code'))
 186
 187
 188 class TestHeapC(TestHeap):
 189     module = c_heapq
 190
 191     def test_comparison_operator(self):
 192         # Issue 3501: Make sure heapq works with both __lt__ and __le__
 193         def hsort(data, comp):
 194             data = map(comp, data)
 195             self.module.heapify(data)
 196             return [self.module.heappop(data).x for i in range(len(data))]
 197         class LT:
 198             def __init__(self, x):
 199                 self.x = x
 200             def __lt__(self, other):
 201                 return self.x > other.x
 202         class LE:
 203             def __init__(self, x):
 204                 self.x = x
 205             def __le__(self, other):
 206                 return self.x >= other.x
 207         data = [random.random() for i in range(100)]
 208         target = sorted(data, reverse=True)
 209         self.assertEqual(hsort(data, LT), target)
 210         self.assertEqual(hsort(data, LE), target)
 211
 212     # As an early adopter, we sanity check the
 213     # test_support.import_fresh_module utility function
 214     def test_accelerated(self):
 215         self.assertTrue(sys.modules['heapq'] is self.module)
 216         self.assertFalse(hasattr(self.module.heapify, 'func_code'))
 217
 218
 219 #==============================================================================
 220
 221 class LenOnly:
 222     "Dummy sequence class defining __len__ but not __getitem__."
 223     def __len__(self):
 224         return 10
 225
 226 class GetOnly:
 227     "Dummy sequence class defining __getitem__ but not __len__."
 228     def __getitem__(self, ndx):
 229         return 10
 230
 231 class CmpErr:
 232     "Dummy element that always raises an error during comparison"
 233     def __cmp__(self, other):
 234         raise ZeroDivisionError
 235
 236 def R(seqn):
 237     'Regular generator'
 238     for i in seqn:
 239         yield i
 240
 241 class G:
 242     'Sequence using __getitem__'
 243     def __init__(self, seqn):
 244         self.seqn = seqn
 245     def __getitem__(self, i):
 246         return self.seqn[i]
 247
 248 class I:
 249     'Sequence using iterator protocol'
 250     def __init__(self, seqn):
 251         self.seqn = seqn
 252         self.i = 0
 253     def __iter__(self):
 254         return self
 255     def next(self):
 256         if self.i >= len(self.seqn): raise StopIteration
 257         v = self.seqn[self.i]
 258         self.i += 1
 259         return v
 260
 261 class Ig:
 262     'Sequence using iterator protocol defined with a generator'
 263     def __init__(self, seqn):
 264         self.seqn = seqn
 265         self.i = 0
 266     def __iter__(self):
 267         for val in self.seqn:
 268             yield val
 269
 270 class X:
 271     'Missing __getitem__ and __iter__'
 272     def __init__(self, seqn):
 273         self.seqn = seqn
 274         self.i = 0
 275     def next(self):
 276         if self.i >= len(self.seqn): raise StopIteration
 277         v = self.seqn[self.i]
 278         self.i += 1
 279         return v
 280
 281 class N:
 282     'Iterator missing next()'
 283     def __init__(self, seqn):
 284         self.seqn = seqn
 285         self.i = 0
 286     def __iter__(self):
 287         return self
 288
 289 class E:
 290     'Test propagation of exceptions'
 291     def __init__(self, seqn):
 292         self.seqn = seqn
 293         self.i = 0
 294     def __iter__(self):
 295         return self
 296     def next(self):
 297         3 // 0
 298
 299 class S:
 300     'Test immediate stop'
 301     def __init__(self, seqn):
 302         pass
 303     def __iter__(self):
 304         return self
 305     def next(self):
 306         raise StopIteration
 307
 308 from itertools import chain, imap
 309 def L(seqn):
 310     'Test multiple tiers of iterators'
 311     return chain(imap(lambda x:x, R(Ig(G(seqn)))))
 312
 313 class TestErrorHandling(unittest.TestCase):
 314     # only for C implementation
 315     module = c_heapq
 316
 317     def test_non_sequence(self):
 318         for f in (self.module.heapify, self.module.heappop):
 319             self.assertRaises(TypeError, f, 10)
 320         for f in (self.module.heappush, self.module.heapreplace,
 321                   self.module.nlargest, self.module.nsmallest):
 322             self.assertRaises(TypeError, f, 10, 10)
 323
 324     def test_len_only(self):
 325         for f in (self.module.heapify, self.module.heappop):
 326             self.assertRaises(TypeError, f, LenOnly())
 327         for f in (self.module.heappush, self.module.heapreplace):
 328             self.assertRaises(TypeError, f, LenOnly(), 10)
 329         for f in (self.module.nlargest, self.module.nsmallest):
 330             self.assertRaises(TypeError, f, 2, LenOnly())
 331
 332     def test_get_only(self):
 333         for f in (self.module.heapify, self.module.heappop):
 334             self.assertRaises(TypeError, f, GetOnly())
 335         for f in (self.module.heappush, self.module.heapreplace):
 336             self.assertRaises(TypeError, f, GetOnly(), 10)
 337         for f in (self.module.nlargest, self.module.nsmallest):
 338             self.assertRaises(TypeError, f, 2, GetOnly())
 339
 340     def test_get_only(self):
 341         seq = [CmpErr(), CmpErr(), CmpErr()]
 342         for f in (self.module.heapify, self.module.heappop):
 343             self.assertRaises(ZeroDivisionError, f, seq)
 344         for f in (self.module.heappush, self.module.heapreplace):
 345             self.assertRaises(ZeroDivisionError, f, seq, 10)
 346         for f in (self.module.nlargest, self.module.nsmallest):
 347             self.assertRaises(ZeroDivisionError, f, 2, seq)
 348
 349     def test_arg_parsing(self):
 350         for f in (self.module.heapify, self.module.heappop,
 351                   self.module.heappush, self.module.heapreplace,
 352                   self.module.nlargest, self.module.nsmallest):
 353             self.assertRaises(TypeError, f, 10)
 354
 355     def test_iterable_args(self):
 356         for f in (self.module.nlargest, self.module.nsmallest):
 357             for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
 358                 for g in (G, I, Ig, L, R):
 359                     with test_support.check_py3k_warnings(
 360                             ("comparing unequal types not supported",
 361                              DeprecationWarning), quiet=True):
 362                         self.assertEqual(f(2, g(s)), f(2,s))
 363                 self.assertEqual(f(2, S(s)), [])
 364                 self.assertRaises(TypeError, f, 2, X(s))
 365                 self.assertRaises(TypeError, f, 2, N(s))
 366                 self.assertRaises(ZeroDivisionError, f, 2, E(s))
 367
 368
 369 #==============================================================================
 370
 371
 372 def test_main(verbose=None):
 373     test_classes = [TestHeapPython, TestHeapC, TestErrorHandling]
 374     test_support.run_unittest(*test_classes)
 375
 376     # verify reference counting
 377     if verbose and hasattr(sys, "gettotalrefcount"):
 378         import gc
 379         counts = [None] * 5
 380         for i in xrange(len(counts)):
 381             test_support.run_unittest(*test_classes)
 382             gc.collect()
 383             counts[i] = sys.gettotalrefcount()
 384         print counts
 385
 386 if __name__ == "__main__":
 387     test_main(verbose=True)