10c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi"""Unittests for heapq.""" 20c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 30c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport sys 40c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiimport random 50c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 60c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom test import test_support 70c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom unittest import TestCase, skipUnless 80c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 90c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yipy_heapq = test_support.import_fresh_module('heapq', blocked=['_heapq']) 100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yic_heapq = test_support.import_fresh_module('heapq', fresh=['_heapq']) 110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi# _heapq.nlargest/nsmallest are saved in heapq._nlargest/_smallest when 130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi# _heapq is imported, so check them there 140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifunc_names = ['heapify', 'heappop', 'heappush', 'heappushpop', 150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'heapreplace', '_nlargest', '_nsmallest'] 160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass TestModules(TestCase): 180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_py_functions(self): 190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for fname in func_names: 200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(getattr(py_heapq, fname).__module__, 'heapq') 210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi @skipUnless(c_heapq, 'requires _heapq') 230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_c_functions(self): 240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for fname in func_names: 250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq') 260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass TestHeap(TestCase): 290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = None 300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_push_pop(self): 320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # 1) Push 256 random numbers and pop them off, verifying all's OK. 330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = [] 340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = [] 350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.check_invariant(heap) 360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for i in range(256): 370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi item = random.random() 380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data.append(item) 390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heappush(heap, item) 400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.check_invariant(heap) 410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi results = [] 420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi while heap: 430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi item = self.module.heappop(heap) 440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.check_invariant(heap) 450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi results.append(item) 460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data_sorted = data[:] 470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data_sorted.sort() 480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(data_sorted, results) 490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # 2) Check that the invariant holds for a sorted array 500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.check_invariant(results) 510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, self.module.heappush, []) 530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi try: 540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, self.module.heappush, None, None) 550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, self.module.heappop, None) 560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi except AttributeError: 570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pass 580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def check_invariant(self, heap): 600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Check the heap invariant. 610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for pos, item in enumerate(heap): 620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if pos: # pos 0 has no parent 630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi parentpos = (pos-1) >> 1 640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertTrue(heap[parentpos] <= item) 650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_heapify(self): 670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for size in range(30): 680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = [random.random() for dummy in range(size)] 690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heapify(heap) 700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.check_invariant(heap) 710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, self.module.heapify, None) 730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_naive_nbest(self): 750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = [random.randrange(2000) for i in range(1000)] 760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = [] 770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for item in data: 780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heappush(heap, item) 790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if len(heap) > 10: 800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heappop(heap) 810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap.sort() 820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(heap, sorted(data)[-10:]) 830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def heapiter(self, heap): 850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # An iterator returning a heap's elements, smallest-first. 860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi try: 870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi while 1: 880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi yield self.module.heappop(heap) 890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi except IndexError: 900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pass 910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_nbest(self): 930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Less-naive "N-best" algorithm, much faster (if len(data) is big 940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # enough <wink>) than sorting all of data. However, if we had a max 950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # heap instead of a min heap, it could go faster still via 960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # heapify'ing all of data (linear time), then doing 10 heappops 970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # (10 log-time steps). 980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = [random.randrange(2000) for i in range(1000)] 990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = data[:10] 1000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heapify(heap) 1010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for item in data[10:]: 1020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if item > heap[0]: # this gets rarer the longer we run 1030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heapreplace(heap, item) 1040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:]) 1050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, self.module.heapreplace, None) 1070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, self.module.heapreplace, None, None) 1080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(IndexError, self.module.heapreplace, [], None) 1090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_nbest_with_pushpop(self): 1110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = [random.randrange(2000) for i in range(1000)] 1120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = data[:10] 1130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heapify(heap) 1140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for item in data[10:]: 1150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heappushpop(heap, item) 1160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:]) 1170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(self.module.heappushpop([], 'x'), 'x') 1180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_heappushpop(self): 1200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi h = [] 1210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi x = self.module.heappushpop(h, 10) 1220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual((h, x), ([], 10)) 1230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi h = [10] 1250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi x = self.module.heappushpop(h, 10.0) 1260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual((h, x), ([10], 10.0)) 1270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(type(h[0]), int) 1280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(type(x), float) 1290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi h = [10]; 1310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi x = self.module.heappushpop(h, 9) 1320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual((h, x), ([10], 9)) 1330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi h = [10]; 1350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi x = self.module.heappushpop(h, 11) 1360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual((h, x), ([11], 10)) 1370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_heapsort(self): 1390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Exercise everything with repeated heapsort checks 1400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for trial in xrange(100): 1410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi size = random.randrange(50) 1420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = [random.randrange(25) for i in range(size)] 1430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if trial & 1: # Half of the time, use heapify 1440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = data[:] 1450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heapify(heap) 1460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi else: # The rest of the time, use heappush 1470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = [] 1480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for item in data: 1490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heappush(heap, item) 1500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap_sorted = [self.module.heappop(heap) for i in range(size)] 1510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(heap_sorted, sorted(data)) 1520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_merge(self): 1540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi inputs = [] 1550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for i in xrange(random.randrange(5)): 1560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi row = sorted(random.randrange(1000) for j in range(random.randrange(10))) 1570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi inputs.append(row) 1580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(sorted(chain(*inputs)), list(self.module.merge(*inputs))) 1590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(list(self.module.merge()), []) 1600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_merge_stability(self): 1620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi class Int(int): 1630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pass 1640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi inputs = [[], [], [], []] 1650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for i in range(20000): 1660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi stream = random.randrange(4) 1670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi x = random.randrange(500) 1680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi obj = Int(x) 1690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi obj.pair = (x, stream) 1700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi inputs[stream].append(obj) 1710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for stream in inputs: 1720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi stream.sort() 1730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi result = [i.pair for i in self.module.merge(*inputs)] 1740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(result, sorted(result)) 1750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_nsmallest(self): 1770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = [(random.randrange(2000), i) for i in range(1000)] 1780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (None, lambda x: x[0] * 547 % 2000): 1790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100): 1800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(self.module.nsmallest(n, data), sorted(data)[:n]) 1810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(self.module.nsmallest(n, data, key=f), 1820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi sorted(data, key=f)[:n]) 1830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_nlargest(self): 1850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = [(random.randrange(2000), i) for i in range(1000)] 1860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (None, lambda x: x[0] * 547 % 2000): 1870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100): 1880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(self.module.nlargest(n, data), 1890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi sorted(data, reverse=True)[:n]) 1900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(self.module.nlargest(n, data, key=f), 1910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi sorted(data, key=f, reverse=True)[:n]) 1920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 1930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_comparison_operator(self): 1940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Issue 3051: Make sure heapq works with both __lt__ and __le__ 1950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def hsort(data, comp): 1960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = map(comp, data) 1970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heapify(data) 1980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return [self.module.heappop(data).x for i in range(len(data))] 1990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi class LT: 2000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, x): 2010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.x = x 2020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __lt__(self, other): 2030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self.x > other.x 2040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi class LE: 2050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, x): 2060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.x = x 2070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __le__(self, other): 2080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self.x >= other.x 2090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi data = [random.random() for i in range(100)] 2100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi target = sorted(data, reverse=True) 2110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(hsort(data, LT), target) 2120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(hsort(data, LE), target) 2130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass TestHeapPython(TestHeap): 2160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = py_heapq 2170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi@skipUnless(c_heapq, 'requires _heapq') 2200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass TestHeapC(TestHeap): 2210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = c_heapq 2220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi#============================================================================== 2250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass LenOnly: 2270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi "Dummy sequence class defining __len__ but not __getitem__." 2280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __len__(self): 2290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return 10 2300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass GetOnly: 2320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi "Dummy sequence class defining __getitem__ but not __len__." 2330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __getitem__(self, ndx): 2340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return 10 2350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass CmpErr: 2370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi "Dummy element that always raises an error during comparison" 2380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __cmp__(self, other): 2390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi raise ZeroDivisionError 2400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef R(seqn): 2420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Regular generator' 2430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for i in seqn: 2440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi yield i 2450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass G: 2470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Sequence using __getitem__' 2480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, seqn): 2490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.seqn = seqn 2500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __getitem__(self, i): 2510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self.seqn[i] 2520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass I: 2540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Sequence using iterator protocol' 2550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, seqn): 2560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.seqn = seqn 2570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.i = 0 2580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __iter__(self): 2590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self 2600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def next(self): 2610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if self.i >= len(self.seqn): raise StopIteration 2620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi v = self.seqn[self.i] 2630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.i += 1 2640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return v 2650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass Ig: 2670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Sequence using iterator protocol defined with a generator' 2680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, seqn): 2690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.seqn = seqn 2700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.i = 0 2710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __iter__(self): 2720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for val in self.seqn: 2730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi yield val 2740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass X: 2760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Missing __getitem__ and __iter__' 2770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, seqn): 2780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.seqn = seqn 2790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.i = 0 2800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def next(self): 2810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if self.i >= len(self.seqn): raise StopIteration 2820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi v = self.seqn[self.i] 2830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.i += 1 2840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return v 2850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass N: 2870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Iterator missing next()' 2880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, seqn): 2890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.seqn = seqn 2900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.i = 0 2910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __iter__(self): 2920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self 2930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 2940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass E: 2950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Test propagation of exceptions' 2960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, seqn): 2970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.seqn = seqn 2980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.i = 0 2990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __iter__(self): 3000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self 3010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def next(self): 3020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3 // 0 3030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass S: 3050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Test immediate stop' 3060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, seqn): 3070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi pass 3080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __iter__(self): 3090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self 3100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def next(self): 3110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi raise StopIteration 3120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yifrom itertools import chain, imap 3140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef L(seqn): 3150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 'Test multiple tiers of iterators' 3160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return chain(imap(lambda x:x, R(Ig(G(seqn))))) 3170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass SideEffectLT: 3190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __init__(self, value, heap): 3200c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.value = value 3210c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.heap = heap 3220c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3230c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def __lt__(self, other): 3240c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.heap[:] = [] 3250c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi return self.value < other.value 3260c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3270c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3280c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass TestErrorHandling(TestCase): 3290c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = None 3300c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3310c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_non_sequence(self): 3320c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.heapify, self.module.heappop): 3330c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises((TypeError, AttributeError), f, 10) 3340c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.heappush, self.module.heapreplace, 3350c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.nlargest, self.module.nsmallest): 3360c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises((TypeError, AttributeError), f, 10, 10) 3370c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3380c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_len_only(self): 3390c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.heapify, self.module.heappop): 3400c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises((TypeError, AttributeError), f, LenOnly()) 3410c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.heappush, self.module.heapreplace): 3420c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises((TypeError, AttributeError), f, LenOnly(), 10) 3430c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.nlargest, self.module.nsmallest): 3440c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, f, 2, LenOnly()) 3450c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3460c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_get_only(self): 3470c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi seq = [CmpErr(), CmpErr(), CmpErr()] 3480c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.heapify, self.module.heappop): 3490c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(ZeroDivisionError, f, seq) 3500c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.heappush, self.module.heapreplace): 3510c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(ZeroDivisionError, f, seq, 10) 3520c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.nlargest, self.module.nsmallest): 3530c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(ZeroDivisionError, f, 2, seq) 3540c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3550c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_arg_parsing(self): 3560c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.heapify, self.module.heappop, 3570c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heappush, self.module.heapreplace, 3580c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.nlargest, self.module.nsmallest): 3590c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises((TypeError, AttributeError), f, 10) 3600c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3610c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_iterable_args(self): 3620c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for f in (self.module.nlargest, self.module.nsmallest): 3630c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): 3640c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for g in (G, I, Ig, L, R): 3650c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with test_support.check_py3k_warnings( 3660c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi ("comparing unequal types not supported", 3670c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi DeprecationWarning), quiet=True): 3680c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(f(2, g(s)), f(2,s)) 3690c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertEqual(f(2, S(s)), []) 3700c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, f, 2, X(s)) 3710c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(TypeError, f, 2, N(s)) 3720c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.assertRaises(ZeroDivisionError, f, 2, E(s)) 3730c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3740c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Issue #17278: the heap may change size while it's being walked. 3750c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3760c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_heappush_mutating_heap(self): 3770c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = [] 3780c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap.extend(SideEffectLT(i, heap) for i in range(200)) 3790c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Python version raises IndexError, C version RuntimeError 3800c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with self.assertRaises((IndexError, RuntimeError)): 3810c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heappush(heap, SideEffectLT(5, heap)) 3820c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3830c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi def test_heappop_mutating_heap(self): 3840c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap = [] 3850c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi heap.extend(SideEffectLT(i, heap) for i in range(200)) 3860c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # Python version raises IndexError, C version RuntimeError 3870c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi with self.assertRaises((IndexError, RuntimeError)): 3880c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi self.module.heappop(heap) 3890c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3900c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3910c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass TestErrorHandlingPython(TestErrorHandling): 3920c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = py_heapq 3930c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3940c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3950c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi@skipUnless(c_heapq, 'requires _heapq') 3960c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiclass TestErrorHandlingC(TestErrorHandling): 3970c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi module = c_heapq 3980c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 3990c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 4000c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi#============================================================================== 4010c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 4020c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 4030c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yidef test_main(verbose=None): 4040c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi test_classes = [TestModules, TestHeapPython, TestHeapC, 4050c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi TestErrorHandlingPython, TestErrorHandlingC] 4060c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi test_support.run_unittest(*test_classes) 4070c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 4080c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi # verify reference counting 4090c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi if verbose and hasattr(sys, "gettotalrefcount"): 4100c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi import gc 4110c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi counts = [None] * 5 4120c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi for i in xrange(len(counts)): 4130c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi test_support.run_unittest(*test_classes) 4140c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi gc.collect() 4150c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi counts[i] = sys.gettotalrefcount() 4160c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi print counts 4170c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi 4180c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yiif __name__ == "__main__": 4190c5958b1636c47ed7c284f859c8e805fd06a0e6Bill Yi test_main(verbose=True) 420