14710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm"""Unittests for heapq.""" 24710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 34710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmimport sys 44710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmimport random 54710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 64710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom test import test_support 74710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom unittest import TestCase, skipUnless 84710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 94710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmpy_heapq = test_support.import_fresh_module('heapq', blocked=['_heapq']) 104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmc_heapq = test_support.import_fresh_module('heapq', fresh=['_heapq']) 114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# _heapq.nlargest/nsmallest are saved in heapq._nlargest/_smallest when 134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# _heapq is imported, so check them there 144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfunc_names = ['heapify', 'heappop', 'heappush', 'heappushpop', 154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'heapreplace', '_nlargest', '_nsmallest'] 164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass TestModules(TestCase): 184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_py_functions(self): 194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for fname in func_names: 204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(getattr(py_heapq, fname).__module__, 'heapq') 214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm @skipUnless(c_heapq, 'requires _heapq') 234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_c_functions(self): 244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for fname in func_names: 254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq') 264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass TestHeap(TestCase): 294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm module = None 304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_push_pop(self): 324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # 1) Push 256 random numbers and pop them off, verifying all's OK. 334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap = [] 344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = [] 354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.check_invariant(heap) 364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for i in range(256): 374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm item = random.random() 384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data.append(item) 394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heappush(heap, item) 404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.check_invariant(heap) 414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm results = [] 424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm while heap: 434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm item = self.module.heappop(heap) 444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.check_invariant(heap) 454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm results.append(item) 464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data_sorted = data[:] 474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data_sorted.sort() 484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(data_sorted, results) 494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # 2) Check that the invariant holds for a sorted array 504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.check_invariant(results) 514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, self.module.heappush, []) 534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm try: 544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, self.module.heappush, None, None) 554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, self.module.heappop, None) 564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm except AttributeError: 574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm pass 584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def check_invariant(self, heap): 604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # Check the heap invariant. 614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for pos, item in enumerate(heap): 624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if pos: # pos 0 has no parent 634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm parentpos = (pos-1) >> 1 644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertTrue(heap[parentpos] <= item) 654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_heapify(self): 674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for size in range(30): 684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap = [random.random() for dummy in range(size)] 694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heapify(heap) 704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.check_invariant(heap) 714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, self.module.heapify, None) 734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_naive_nbest(self): 754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = [random.randrange(2000) for i in range(1000)] 764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap = [] 774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for item in data: 784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heappush(heap, item) 794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if len(heap) > 10: 804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heappop(heap) 814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap.sort() 824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(heap, sorted(data)[-10:]) 834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def heapiter(self, heap): 854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # An iterator returning a heap's elements, smallest-first. 864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm try: 874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm while 1: 884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm yield self.module.heappop(heap) 894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm except IndexError: 904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm pass 914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_nbest(self): 934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # Less-naive "N-best" algorithm, much faster (if len(data) is big 944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # enough <wink>) than sorting all of data. However, if we had a max 954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # heap instead of a min heap, it could go faster still via 964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # heapify'ing all of data (linear time), then doing 10 heappops 974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # (10 log-time steps). 984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = [random.randrange(2000) for i in range(1000)] 994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap = data[:10] 1004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heapify(heap) 1014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for item in data[10:]: 1024710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if item > heap[0]: # this gets rarer the longer we run 1034710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heapreplace(heap, item) 1044710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:]) 1054710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1064710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, self.module.heapreplace, None) 1074710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, self.module.heapreplace, None, None) 1084710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(IndexError, self.module.heapreplace, [], None) 1094710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_nbest_with_pushpop(self): 1114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = [random.randrange(2000) for i in range(1000)] 1124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap = data[:10] 1134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heapify(heap) 1144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for item in data[10:]: 1154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heappushpop(heap, item) 1164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:]) 1174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(self.module.heappushpop([], 'x'), 'x') 1184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_heappushpop(self): 1204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm h = [] 1214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm x = self.module.heappushpop(h, 10) 1224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual((h, x), ([], 10)) 1234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm h = [10] 1254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm x = self.module.heappushpop(h, 10.0) 1264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual((h, x), ([10], 10.0)) 1274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(type(h[0]), int) 1284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(type(x), float) 1294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm h = [10]; 1314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm x = self.module.heappushpop(h, 9) 1324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual((h, x), ([10], 9)) 1334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm h = [10]; 1354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm x = self.module.heappushpop(h, 11) 1364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual((h, x), ([11], 10)) 1374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_heapsort(self): 1394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # Exercise everything with repeated heapsort checks 1404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for trial in xrange(100): 1414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm size = random.randrange(50) 1424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = [random.randrange(25) for i in range(size)] 1434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if trial & 1: # Half of the time, use heapify 1444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap = data[:] 1454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heapify(heap) 1464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm else: # The rest of the time, use heappush 1474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap = [] 1484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for item in data: 1494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heappush(heap, item) 1504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm heap_sorted = [self.module.heappop(heap) for i in range(size)] 1514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(heap_sorted, sorted(data)) 1524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_merge(self): 1544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm inputs = [] 1554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for i in xrange(random.randrange(5)): 1564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm row = sorted(random.randrange(1000) for j in range(random.randrange(10))) 1574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm inputs.append(row) 1584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(sorted(chain(*inputs)), list(self.module.merge(*inputs))) 1594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(list(self.module.merge()), []) 1604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_merge_stability(self): 1624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm class Int(int): 1634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm pass 1644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm inputs = [[], [], [], []] 1654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for i in range(20000): 1664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm stream = random.randrange(4) 1674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm x = random.randrange(500) 1684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm obj = Int(x) 1694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm obj.pair = (x, stream) 1704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm inputs[stream].append(obj) 1714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for stream in inputs: 1724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm stream.sort() 1734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm result = [i.pair for i in self.module.merge(*inputs)] 1744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(result, sorted(result)) 1754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_nsmallest(self): 1774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = [(random.randrange(2000), i) for i in range(1000)] 1784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (None, lambda x: x[0] * 547 % 2000): 1794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100): 1804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(self.module.nsmallest(n, data), sorted(data)[:n]) 1814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(self.module.nsmallest(n, data, key=f), 1824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm sorted(data, key=f)[:n]) 1834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_nlargest(self): 1854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = [(random.randrange(2000), i) for i in range(1000)] 1864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (None, lambda x: x[0] * 547 % 2000): 1874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100): 1884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(self.module.nlargest(n, data), 1894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm sorted(data, reverse=True)[:n]) 1904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(self.module.nlargest(n, data, key=f), 1914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm sorted(data, key=f, reverse=True)[:n]) 1924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_comparison_operator(self): 1944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # Issue 3051: Make sure heapq works with both __lt__ and __le__ 1954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def hsort(data, comp): 1964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = map(comp, data) 1974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heapify(data) 1984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return [self.module.heappop(data).x for i in range(len(data))] 1994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm class LT: 2004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, x): 2014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.x = x 2024710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __lt__(self, other): 2034710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return self.x > other.x 2044710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm class LE: 2054710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, x): 2064710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.x = x 2074710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __le__(self, other): 2084710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return self.x >= other.x 2094710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm data = [random.random() for i in range(100)] 2104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm target = sorted(data, reverse=True) 2114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(hsort(data, LT), target) 2124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(hsort(data, LE), target) 2134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass TestHeapPython(TestHeap): 2164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm module = py_heapq 2174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm@skipUnless(c_heapq, 'requires _heapq') 2204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass TestHeapC(TestHeap): 2214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm module = c_heapq 2224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm#============================================================================== 2254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass LenOnly: 2274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm "Dummy sequence class defining __len__ but not __getitem__." 2284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __len__(self): 2294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return 10 2304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass GetOnly: 2324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm "Dummy sequence class defining __getitem__ but not __len__." 2334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __getitem__(self, ndx): 2344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return 10 2354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass CmpErr: 2374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm "Dummy element that always raises an error during comparison" 2384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __cmp__(self, other): 2394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm raise ZeroDivisionError 2404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef R(seqn): 2424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Regular generator' 2434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for i in seqn: 2444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm yield i 2454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass G: 2474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Sequence using __getitem__' 2484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, seqn): 2494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.seqn = seqn 2504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __getitem__(self, i): 2514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return self.seqn[i] 2524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass I: 2544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Sequence using iterator protocol' 2554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, seqn): 2564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.seqn = seqn 2574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.i = 0 2584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __iter__(self): 2594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return self 2604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def next(self): 2614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if self.i >= len(self.seqn): raise StopIteration 2624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm v = self.seqn[self.i] 2634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.i += 1 2644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return v 2654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass Ig: 2674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Sequence using iterator protocol defined with a generator' 2684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, seqn): 2694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.seqn = seqn 2704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.i = 0 2714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __iter__(self): 2724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for val in self.seqn: 2734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm yield val 2744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass X: 2764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Missing __getitem__ and __iter__' 2774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, seqn): 2784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.seqn = seqn 2794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.i = 0 2804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def next(self): 2814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if self.i >= len(self.seqn): raise StopIteration 2824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm v = self.seqn[self.i] 2834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.i += 1 2844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return v 2854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass N: 2874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Iterator missing next()' 2884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, seqn): 2894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.seqn = seqn 2904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.i = 0 2914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __iter__(self): 2924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return self 2934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass E: 2954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Test propagation of exceptions' 2964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, seqn): 2974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.seqn = seqn 2984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.i = 0 2994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __iter__(self): 3004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return self 3014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def next(self): 3024710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3 // 0 3034710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3044710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass S: 3054710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Test immediate stop' 3064710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __init__(self, seqn): 3074710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm pass 3084710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def __iter__(self): 3094710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return self 3104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def next(self): 3114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm raise StopIteration 3124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom itertools import chain, imap 3144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef L(seqn): 3154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'Test multiple tiers of iterators' 3164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return chain(imap(lambda x:x, R(Ig(G(seqn))))) 3174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass TestErrorHandling(TestCase): 3194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm module = None 3204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_non_sequence(self): 3224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.heapify, self.module.heappop): 3234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises((TypeError, AttributeError), f, 10) 3244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.heappush, self.module.heapreplace, 3254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.nlargest, self.module.nsmallest): 3264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises((TypeError, AttributeError), f, 10, 10) 3274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_len_only(self): 3294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.heapify, self.module.heappop): 3304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises((TypeError, AttributeError), f, LenOnly()) 3314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.heappush, self.module.heapreplace): 3324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises((TypeError, AttributeError), f, LenOnly(), 10) 3334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.nlargest, self.module.nsmallest): 3344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, f, 2, LenOnly()) 3354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_get_only(self): 3374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm seq = [CmpErr(), CmpErr(), CmpErr()] 3384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.heapify, self.module.heappop): 3394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(ZeroDivisionError, f, seq) 3404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.heappush, self.module.heapreplace): 3414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(ZeroDivisionError, f, seq, 10) 3424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.nlargest, self.module.nsmallest): 3434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(ZeroDivisionError, f, 2, seq) 3444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_arg_parsing(self): 3464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.heapify, self.module.heappop, 3474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.heappush, self.module.heapreplace, 3484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.module.nlargest, self.module.nsmallest): 3494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises((TypeError, AttributeError), f, 10) 3504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def test_iterable_args(self): 3524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for f in (self.module.nlargest, self.module.nsmallest): 3534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): 3544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for g in (G, I, Ig, L, R): 3554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm with test_support.check_py3k_warnings( 3564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ("comparing unequal types not supported", 3574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm DeprecationWarning), quiet=True): 3584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(f(2, g(s)), f(2,s)) 3594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertEqual(f(2, S(s)), []) 3604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, f, 2, X(s)) 3614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(TypeError, f, 2, N(s)) 3624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm self.assertRaises(ZeroDivisionError, f, 2, E(s)) 3634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass TestErrorHandlingPython(TestErrorHandling): 3664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm module = py_heapq 3674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm@skipUnless(c_heapq, 'requires _heapq') 3704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmclass TestErrorHandlingC(TestErrorHandling): 3714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm module = c_heapq 3724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm#============================================================================== 3754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef test_main(verbose=None): 3784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm test_classes = [TestModules, TestHeapPython, TestHeapC, 3794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm TestErrorHandlingPython, TestErrorHandlingC] 3804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm test_support.run_unittest(*test_classes) 3814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # verify reference counting 3834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if verbose and hasattr(sys, "gettotalrefcount"): 3844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm import gc 3854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm counts = [None] * 5 3864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for i in xrange(len(counts)): 3874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm test_support.run_unittest(*test_classes) 3884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm gc.collect() 3894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm counts[i] = sys.gettotalrefcount() 3904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm print counts 3914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmif __name__ == "__main__": 3934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm test_main(verbose=True) 394