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