blob: 32958696689ede3211df36b1d2e6845992c82117 [file] [log] [blame]
Guido van Rossum0b191782002-08-02 18:29:53 +00001"""Unittests for heapq."""
2
Guido van Rossum0b191782002-08-02 18:29:53 +00003import random
Raymond Hettingerbce036b2004-06-10 05:07:18 +00004import unittest
Benjamin Petersonee8712c2008-05-20 21:35:26 +00005from test import support
Raymond Hettinger2e3dfaf2004-06-13 05:26:33 +00006import sys
Guido van Rossum0b191782002-08-02 18:29:53 +00007
Christian Heimesd3eb5a152008-02-24 00:38:49 +00008# We do a bit of trickery here to be able to test both the C implementation
9# and the Python implementation of the module.
Christian Heimesd3eb5a152008-02-24 00:38:49 +000010import heapq as c_heapq
Nick Coghlan47384702009-04-22 16:13:36 +000011py_heapq = support.import_fresh_module('heapq', blocked=['_heapq'])
Tim Petersaa7d2432002-08-03 02:11:26 +000012
Raymond Hettingerbce036b2004-06-10 05:07:18 +000013class TestHeap(unittest.TestCase):
Christian Heimesd3eb5a152008-02-24 00:38:49 +000014 module = None
Tim Petersaa7d2432002-08-03 02:11:26 +000015
Raymond Hettingerbce036b2004-06-10 05:07:18 +000016 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)
Christian Heimesd3eb5a152008-02-24 00:38:49 +000024 self.module.heappush(heap, item)
Raymond Hettingerbce036b2004-06-10 05:07:18 +000025 self.check_invariant(heap)
26 results = []
27 while heap:
Christian Heimesd3eb5a152008-02-24 00:38:49 +000028 item = self.module.heappop(heap)
Raymond Hettingerbce036b2004-06-10 05:07:18 +000029 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
Christian Heimesd3eb5a152008-02-24 00:38:49 +000037 self.assertRaises(TypeError, self.module.heappush, [])
Raymond Hettingere1defa42004-11-29 05:54:48 +000038 try:
Christian Heimesd3eb5a152008-02-24 00:38:49 +000039 self.assertRaises(TypeError, self.module.heappush, None, None)
40 self.assertRaises(TypeError, self.module.heappop, None)
Raymond Hettingere1defa42004-11-29 05:54:48 +000041 except AttributeError:
42 pass
Neal Norwitzd7be1182004-07-08 01:56:46 +000043
Raymond Hettingerbce036b2004-06-10 05:07:18 +000044 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
Benjamin Petersonc9c0f202009-06-30 23:06:06 +000049 self.assertTrue(heap[parentpos] <= item)
Raymond Hettingerbce036b2004-06-10 05:07:18 +000050
51 def test_heapify(self):
52 for size in range(30):
53 heap = [random.random() for dummy in range(size)]
Christian Heimesd3eb5a152008-02-24 00:38:49 +000054 self.module.heapify(heap)
Raymond Hettingerbce036b2004-06-10 05:07:18 +000055 self.check_invariant(heap)
56
Christian Heimesd3eb5a152008-02-24 00:38:49 +000057 self.assertRaises(TypeError, self.module.heapify, None)
Neal Norwitzd7be1182004-07-08 01:56:46 +000058
Raymond Hettingerbce036b2004-06-10 05:07:18 +000059 def test_naive_nbest(self):
60 data = [random.randrange(2000) for i in range(1000)]
61 heap = []
62 for item in data:
Christian Heimesd3eb5a152008-02-24 00:38:49 +000063 self.module.heappush(heap, item)
Raymond Hettingerbce036b2004-06-10 05:07:18 +000064 if len(heap) > 10:
Christian Heimesd3eb5a152008-02-24 00:38:49 +000065 self.module.heappop(heap)
Raymond Hettingerbce036b2004-06-10 05:07:18 +000066 heap.sort()
67 self.assertEqual(heap, sorted(data)[-10:])
68
Christian Heimesd3eb5a152008-02-24 00:38:49 +000069 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
Raymond Hettingerbce036b2004-06-10 05:07:18 +000077 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]
Christian Heimesd3eb5a152008-02-24 00:38:49 +000085 self.module.heapify(heap)
Raymond Hettingerbce036b2004-06-10 05:07:18 +000086 for item in data[10:]:
87 if item > heap[0]: # this gets rarer the longer we run
Christian Heimesd3eb5a152008-02-24 00:38:49 +000088 self.module.heapreplace(heap, item)
89 self.assertEqual(list(self.heapiter(heap)), sorted(data)[-10:])
Raymond Hettingerbce036b2004-06-10 05:07:18 +000090
Christian Heimesd3eb5a152008-02-24 00:38:49 +000091 self.assertRaises(TypeError, self.module.heapreplace, None)
92 self.assertRaises(TypeError, self.module.heapreplace, None, None)
93 self.assertRaises(IndexError, self.module.heapreplace, [], None)
Neal Norwitzd7be1182004-07-08 01:56:46 +000094
Christian Heimesdd15f6c2008-03-16 00:07:10 +000095 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
Raymond Hettingerbce036b2004-06-10 05:07:18 +0000123 def test_heapsort(self):
124 # Exercise everything with repeated heapsort checks
Guido van Rossum805365e2007-05-07 22:24:25 +0000125 for trial in range(100):
Raymond Hettingerbce036b2004-06-10 05:07:18 +0000126 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[:]
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000130 self.module.heapify(heap)
Raymond Hettingerbce036b2004-06-10 05:07:18 +0000131 else: # The rest of the time, use heappush
132 heap = []
133 for item in data:
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000134 self.module.heappush(heap, item)
135 heap_sorted = [self.module.heappop(heap) for i in range(size)]
Raymond Hettingerbce036b2004-06-10 05:07:18 +0000136 self.assertEqual(heap_sorted, sorted(data))
137
Thomas Wouterscf297e42007-02-23 15:07:44 +0000138 def test_merge(self):
139 inputs = []
Guido van Rossum805365e2007-05-07 22:24:25 +0000140 for i in range(random.randrange(5)):
Thomas Wouterscf297e42007-02-23 15:07:44 +0000141 row = sorted(random.randrange(1000) for j in range(random.randrange(10)))
142 inputs.append(row)
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000143 self.assertEqual(sorted(chain(*inputs)), list(self.module.merge(*inputs)))
144 self.assertEqual(list(self.module.merge()), [])
Thomas Wouterscf297e42007-02-23 15:07:44 +0000145
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()
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000158 result = [i.pair for i in self.module.merge(*inputs)]
Thomas Wouterscf297e42007-02-23 15:07:44 +0000159 self.assertEqual(result, sorted(result))
160
Raymond Hettingerbce036b2004-06-10 05:07:18 +0000161 def test_nsmallest(self):
Thomas Wouters902d6eb2007-01-09 23:18:33 +0000162 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):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000165 self.assertEqual(list(self.module.nsmallest(n, data)),
166 sorted(data)[:n])
167 self.assertEqual(list(self.module.nsmallest(n, data, key=f)),
Thomas Wouters902d6eb2007-01-09 23:18:33 +0000168 sorted(data, key=f)[:n])
Raymond Hettingerbce036b2004-06-10 05:07:18 +0000169
Raymond Hettinger4901a1f2004-12-02 08:59:14 +0000170 def test_nlargest(self):
Thomas Wouters902d6eb2007-01-09 23:18:33 +0000171 data = [(random.randrange(2000), i) for i in range(1000)]
172 for f in (None, lambda x: x[0] * 547 % 2000):
173 for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000174 self.assertEqual(list(self.module.nlargest(n, data)),
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000175 sorted(data, reverse=True)[:n])
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000176 self.assertEqual(list(self.module.nlargest(n, data, key=f)),
Thomas Wouters902d6eb2007-01-09 23:18:33 +0000177 sorted(data, key=f, reverse=True)[:n])
Tim Petersaa7d2432002-08-03 02:11:26 +0000178
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000179class TestHeapPython(TestHeap):
180 module = py_heapq
181
Nick Coghlanfce769e2009-04-11 14:30:59 +0000182 # As an early adopter, we sanity check the
183 # test.support.import_fresh_module utility function
184 def test_pure_python(self):
185 self.assertFalse(sys.modules['heapq'] is self.module)
186 self.assertTrue(hasattr(self.module.heapify, '__code__'))
187
188
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000189class TestHeapC(TestHeap):
190 module = c_heapq
191
Amaury Forgeot d'Arcfdfe62d2008-06-17 20:36:03 +0000192 def test_comparison_operator(self):
Amaury Forgeot d'Arc2ba198d2008-06-17 21:25:35 +0000193 # Issue 3501: Make sure heapq works with both __lt__
194 # For python 3.0, __le__ alone is not enough
Amaury Forgeot d'Arcfdfe62d2008-06-17 20:36:03 +0000195 def hsort(data, comp):
196 data = [comp(x) for x in data]
197 self.module.heapify(data)
198 return [self.module.heappop(data).x for i in range(len(data))]
199 class LT:
200 def __init__(self, x):
201 self.x = x
202 def __lt__(self, other):
203 return self.x > other.x
204 class LE:
205 def __init__(self, x):
206 self.x = x
Amaury Forgeot d'Arc35c86582008-06-17 21:11:29 +0000207 def __le__(self, other):
Amaury Forgeot d'Arcfdfe62d2008-06-17 20:36:03 +0000208 return self.x >= other.x
209 data = [random.random() for i in range(100)]
210 target = sorted(data, reverse=True)
211 self.assertEqual(hsort(data, LT), target)
Amaury Forgeot d'Arc2ba198d2008-06-17 21:25:35 +0000212 self.assertRaises(TypeError, data, LE)
Amaury Forgeot d'Arcfdfe62d2008-06-17 20:36:03 +0000213
Nick Coghlanfce769e2009-04-11 14:30:59 +0000214 # As an early adopter, we sanity check the
215 # test.support.import_fresh_module utility function
216 def test_accelerated(self):
217 self.assertTrue(sys.modules['heapq'] is self.module)
218 self.assertFalse(hasattr(self.module.heapify, '__code__'))
219
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000220
221#==============================================================================
222
223class LenOnly:
224 "Dummy sequence class defining __len__ but not __getitem__."
225 def __len__(self):
226 return 10
227
228class GetOnly:
229 "Dummy sequence class defining __getitem__ but not __len__."
230 def __getitem__(self, ndx):
231 return 10
232
233class CmpErr:
234 "Dummy element that always raises an error during comparison"
Mark Dickinsona56c4672009-01-27 18:17:45 +0000235 def __eq__(self, other):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000236 raise ZeroDivisionError
Mark Dickinsona56c4672009-01-27 18:17:45 +0000237 __ne__ = __lt__ = __le__ = __gt__ = __ge__ = __eq__
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000238
239def R(seqn):
240 'Regular generator'
241 for i in seqn:
242 yield i
243
244class G:
245 'Sequence using __getitem__'
246 def __init__(self, seqn):
247 self.seqn = seqn
248 def __getitem__(self, i):
249 return self.seqn[i]
250
251class I:
252 'Sequence using iterator protocol'
253 def __init__(self, seqn):
254 self.seqn = seqn
255 self.i = 0
256 def __iter__(self):
257 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000258 def __next__(self):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000259 if self.i >= len(self.seqn): raise StopIteration
260 v = self.seqn[self.i]
261 self.i += 1
262 return v
263
264class Ig:
265 'Sequence using iterator protocol defined with a generator'
266 def __init__(self, seqn):
267 self.seqn = seqn
268 self.i = 0
269 def __iter__(self):
270 for val in self.seqn:
271 yield val
272
273class X:
274 'Missing __getitem__ and __iter__'
275 def __init__(self, seqn):
276 self.seqn = seqn
277 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +0000278 def __next__(self):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000279 if self.i >= len(self.seqn): raise StopIteration
280 v = self.seqn[self.i]
281 self.i += 1
282 return v
283
284class N:
Georg Brandla18af4e2007-04-21 15:47:16 +0000285 'Iterator missing __next__()'
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000286 def __init__(self, seqn):
287 self.seqn = seqn
288 self.i = 0
289 def __iter__(self):
290 return self
291
292class E:
293 'Test propagation of exceptions'
294 def __init__(self, seqn):
295 self.seqn = seqn
296 self.i = 0
297 def __iter__(self):
298 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000299 def __next__(self):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000300 3 // 0
301
302class S:
303 'Test immediate stop'
304 def __init__(self, seqn):
305 pass
306 def __iter__(self):
307 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000308 def __next__(self):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000309 raise StopIteration
310
Raymond Hettinger736c0ab2008-03-13 02:09:15 +0000311from itertools import chain
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000312def L(seqn):
313 'Test multiple tiers of iterators'
Raymond Hettingera6c60372008-03-13 01:26:19 +0000314 return chain(map(lambda x:x, R(Ig(G(seqn)))))
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000315
316class TestErrorHandling(unittest.TestCase):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000317 # only for C implementation
318 module = c_heapq
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000319
320 def test_non_sequence(self):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000321 for f in (self.module.heapify, self.module.heappop):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000322 self.assertRaises(TypeError, f, 10)
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000323 for f in (self.module.heappush, self.module.heapreplace,
324 self.module.nlargest, self.module.nsmallest):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000325 self.assertRaises(TypeError, f, 10, 10)
326
327 def test_len_only(self):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000328 for f in (self.module.heapify, self.module.heappop):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000329 self.assertRaises(TypeError, f, LenOnly())
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000330 for f in (self.module.heappush, self.module.heapreplace):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000331 self.assertRaises(TypeError, f, LenOnly(), 10)
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000332 for f in (self.module.nlargest, self.module.nsmallest):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000333 self.assertRaises(TypeError, f, 2, LenOnly())
334
335 def test_get_only(self):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000336 for f in (self.module.heapify, self.module.heappop):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000337 self.assertRaises(TypeError, f, GetOnly())
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000338 for f in (self.module.heappush, self.module.heapreplace):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000339 self.assertRaises(TypeError, f, GetOnly(), 10)
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000340 for f in (self.module.nlargest, self.module.nsmallest):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000341 self.assertRaises(TypeError, f, 2, GetOnly())
342
343 def test_get_only(self):
344 seq = [CmpErr(), CmpErr(), CmpErr()]
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000345 for f in (self.module.heapify, self.module.heappop):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000346 self.assertRaises(ZeroDivisionError, f, seq)
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000347 for f in (self.module.heappush, self.module.heapreplace):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000348 self.assertRaises(ZeroDivisionError, f, seq, 10)
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000349 for f in (self.module.nlargest, self.module.nsmallest):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000350 self.assertRaises(ZeroDivisionError, f, 2, seq)
351
352 def test_arg_parsing(self):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000353 for f in (self.module.heapify, self.module.heappop,
354 self.module.heappush, self.module.heapreplace,
355 self.module.nlargest, self.module.nsmallest):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000356 self.assertRaises(TypeError, f, 10)
357
358 def test_iterable_args(self):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000359 for f in (self.module.nlargest, self.module.nsmallest):
Guido van Rossum805365e2007-05-07 22:24:25 +0000360 for s in ("123", "", range(1000), (1, 1.2), range(2000,2200,5)):
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000361 for g in (G, I, Ig, L, R):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000362 self.assertEqual(list(f(2, g(s))), list(f(2,s)))
363 self.assertEqual(list(f(2, S(s))), [])
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000364 self.assertRaises(TypeError, f, 2, X(s))
365 self.assertRaises(TypeError, f, 2, N(s))
366 self.assertRaises(ZeroDivisionError, f, 2, E(s))
367
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000368
Raymond Hettinger855d9a92004-09-28 00:03:54 +0000369#==============================================================================
370
371
Raymond Hettinger2e3dfaf2004-06-13 05:26:33 +0000372def test_main(verbose=None):
Christian Heimesd3eb5a152008-02-24 00:38:49 +0000373 test_classes = [TestHeapPython, TestHeapC, TestErrorHandling]
Benjamin Petersonee8712c2008-05-20 21:35:26 +0000374 support.run_unittest(*test_classes)
Raymond Hettinger2e3dfaf2004-06-13 05:26:33 +0000375
376 # verify reference counting
377 if verbose and hasattr(sys, "gettotalrefcount"):
378 import gc
379 counts = [None] * 5
Guido van Rossum805365e2007-05-07 22:24:25 +0000380 for i in range(len(counts)):
Benjamin Petersonee8712c2008-05-20 21:35:26 +0000381 support.run_unittest(*test_classes)
Raymond Hettinger2e3dfaf2004-06-13 05:26:33 +0000382 gc.collect()
383 counts[i] = sys.gettotalrefcount()
Guido van Rossumbe19ed72007-02-09 05:37:30 +0000384 print(counts)
Guido van Rossum0b191782002-08-02 18:29:53 +0000385
386if __name__ == "__main__":
Raymond Hettinger2e3dfaf2004-06-13 05:26:33 +0000387 test_main(verbose=True)