blob: 56e21f8aa70064d554c3593c4c8a950c992b5204 [file] [log] [blame]
Guido van Rossum8b48cf92001-04-21 13:33:54 +00001# Test iterators.
2
Serhiy Storchaka4faf5c52015-05-21 20:50:25 +03003import sys
Guido van Rossum8b48cf92001-04-21 13:33:54 +00004import unittest
Benjamin Peterson945c5792010-06-22 20:34:34 +00005from test.support import run_unittest, TESTFN, unlink, cpython_only
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +00006import pickle
7import collections.abc
Guido van Rossum8b48cf92001-04-21 13:33:54 +00008
9# Test result of triple loop (too big to inline)
10TRIPLETS = [(0, 0, 0), (0, 0, 1), (0, 0, 2),
11 (0, 1, 0), (0, 1, 1), (0, 1, 2),
12 (0, 2, 0), (0, 2, 1), (0, 2, 2),
13
14 (1, 0, 0), (1, 0, 1), (1, 0, 2),
15 (1, 1, 0), (1, 1, 1), (1, 1, 2),
16 (1, 2, 0), (1, 2, 1), (1, 2, 2),
17
18 (2, 0, 0), (2, 0, 1), (2, 0, 2),
19 (2, 1, 0), (2, 1, 1), (2, 1, 2),
20 (2, 2, 0), (2, 2, 1), (2, 2, 2)]
21
22# Helper classes
23
24class BasicIterClass:
25 def __init__(self, n):
26 self.n = n
27 self.i = 0
Georg Brandla18af4e2007-04-21 15:47:16 +000028 def __next__(self):
Guido van Rossum8b48cf92001-04-21 13:33:54 +000029 res = self.i
30 if res >= self.n:
31 raise StopIteration
32 self.i = res + 1
33 return res
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000034 def __iter__(self):
35 return self
Guido van Rossum8b48cf92001-04-21 13:33:54 +000036
37class IteratingSequenceClass:
38 def __init__(self, n):
39 self.n = n
40 def __iter__(self):
41 return BasicIterClass(self.n)
42
43class SequenceClass:
44 def __init__(self, n):
45 self.n = n
46 def __getitem__(self, i):
47 if 0 <= i < self.n:
48 return i
49 else:
50 raise IndexError
51
Serhiy Storchaka4faf5c52015-05-21 20:50:25 +030052class UnlimitedSequenceClass:
53 def __getitem__(self, i):
54 return i
55
Guido van Rossum8b48cf92001-04-21 13:33:54 +000056# Main test suite
57
58class TestCase(unittest.TestCase):
59
60 # Helper to check that an iterator returns a given sequence
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000061 def check_iterator(self, it, seq, pickle=True):
62 if pickle:
63 self.check_pickle(it, seq)
Guido van Rossum8b48cf92001-04-21 13:33:54 +000064 res = []
65 while 1:
66 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000067 val = next(it)
Guido van Rossum8b48cf92001-04-21 13:33:54 +000068 except StopIteration:
69 break
70 res.append(val)
71 self.assertEqual(res, seq)
72
73 # Helper to check that a for loop generates a given sequence
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000074 def check_for_loop(self, expr, seq, pickle=True):
75 if pickle:
76 self.check_pickle(iter(expr), seq)
Guido van Rossum8b48cf92001-04-21 13:33:54 +000077 res = []
78 for val in expr:
79 res.append(val)
80 self.assertEqual(res, seq)
81
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000082 # Helper to check picklability
83 def check_pickle(self, itorg, seq):
Serhiy Storchakabad12572014-12-15 14:03:42 +020084 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
85 d = pickle.dumps(itorg, proto)
86 it = pickle.loads(d)
87 # Cannot assert type equality because dict iterators unpickle as list
88 # iterators.
89 # self.assertEqual(type(itorg), type(it))
90 self.assertTrue(isinstance(it, collections.abc.Iterator))
91 self.assertEqual(list(it), seq)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +000092
Serhiy Storchakabad12572014-12-15 14:03:42 +020093 it = pickle.loads(d)
94 try:
95 next(it)
96 except StopIteration:
97 continue
98 d = pickle.dumps(it, proto)
99 it = pickle.loads(d)
100 self.assertEqual(list(it), seq[1:])
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000101
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000102 # Test basic use of iter() function
103 def test_iter_basic(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000104 self.check_iterator(iter(range(10)), list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000105
106 # Test that iter(iter(x)) is the same as iter(x)
107 def test_iter_idempotency(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000108 seq = list(range(10))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000109 it = iter(seq)
110 it2 = iter(it)
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000111 self.assertTrue(it is it2)
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000112
113 # Test that for loops over iterators work
114 def test_iter_for_loop(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000115 self.check_for_loop(iter(range(10)), list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000116
117 # Test several independent iterators over the same list
118 def test_iter_independence(self):
119 seq = range(3)
120 res = []
121 for i in iter(seq):
122 for j in iter(seq):
123 for k in iter(seq):
124 res.append((i, j, k))
125 self.assertEqual(res, TRIPLETS)
126
127 # Test triple list comprehension using iterators
128 def test_nested_comprehensions_iter(self):
129 seq = range(3)
130 res = [(i, j, k)
131 for i in iter(seq) for j in iter(seq) for k in iter(seq)]
132 self.assertEqual(res, TRIPLETS)
133
134 # Test triple list comprehension without iterators
135 def test_nested_comprehensions_for(self):
136 seq = range(3)
137 res = [(i, j, k) for i in seq for j in seq for k in seq]
138 self.assertEqual(res, TRIPLETS)
139
140 # Test a class with __iter__ in a for loop
141 def test_iter_class_for(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000142 self.check_for_loop(IteratingSequenceClass(10), list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000143
144 # Test a class with __iter__ with explicit iter()
145 def test_iter_class_iter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000146 self.check_iterator(iter(IteratingSequenceClass(10)), list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000147
148 # Test for loop on a sequence class without __iter__
149 def test_seq_class_for(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000150 self.check_for_loop(SequenceClass(10), list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000151
152 # Test iter() on a sequence class without __iter__
153 def test_seq_class_iter(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000154 self.check_iterator(iter(SequenceClass(10)), list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000155
Serhiy Storchakaaabafe72016-03-06 14:10:24 +0200156 def test_mutating_seq_class_iter_pickle(self):
157 orig = SequenceClass(5)
158 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
159 # initial iterator
160 itorig = iter(orig)
161 d = pickle.dumps((itorig, orig), proto)
162 it, seq = pickle.loads(d)
163 seq.n = 7
164 self.assertIs(type(it), type(itorig))
165 self.assertEqual(list(it), list(range(7)))
166
167 # running iterator
168 next(itorig)
169 d = pickle.dumps((itorig, orig), proto)
170 it, seq = pickle.loads(d)
171 seq.n = 7
172 self.assertIs(type(it), type(itorig))
173 self.assertEqual(list(it), list(range(1, 7)))
174
175 # empty iterator
176 for i in range(1, 5):
177 next(itorig)
178 d = pickle.dumps((itorig, orig), proto)
179 it, seq = pickle.loads(d)
180 seq.n = 7
181 self.assertIs(type(it), type(itorig))
182 self.assertEqual(list(it), list(range(5, 7)))
183
184 # exhausted iterator
185 self.assertRaises(StopIteration, next, itorig)
186 d = pickle.dumps((itorig, orig), proto)
187 it, seq = pickle.loads(d)
188 seq.n = 7
189 self.assertTrue(isinstance(it, collections.abc.Iterator))
190 self.assertEqual(list(it), [])
191
Amaury Forgeot d'Arcf343e012009-01-12 23:58:21 +0000192 # Test a new_style class with __iter__ but no next() method
193 def test_new_style_iter_class(self):
194 class IterClass(object):
195 def __iter__(self):
196 return self
197 self.assertRaises(TypeError, iter, IterClass())
198
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000199 # Test two-argument iter() with callable instance
200 def test_iter_callable(self):
201 class C:
202 def __init__(self):
203 self.i = 0
204 def __call__(self):
205 i = self.i
206 self.i = i + 1
207 if i > 100:
208 raise IndexError # Emergency stop
209 return i
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000210 self.check_iterator(iter(C(), 10), list(range(10)), pickle=False)
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000211
212 # Test two-argument iter() with function
213 def test_iter_function(self):
214 def spam(state=[0]):
215 i = state[0]
216 state[0] = i+1
217 return i
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000218 self.check_iterator(iter(spam, 10), list(range(10)), pickle=False)
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000219
220 # Test two-argument iter() with function that raises StopIteration
221 def test_iter_function_stop(self):
222 def spam(state=[0]):
223 i = state[0]
224 if i == 10:
225 raise StopIteration
226 state[0] = i+1
227 return i
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000228 self.check_iterator(iter(spam, 20), list(range(10)), pickle=False)
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000229
230 # Test exception propagation through function iterator
231 def test_exception_function(self):
232 def spam(state=[0]):
233 i = state[0]
234 state[0] = i+1
235 if i == 10:
236 raise RuntimeError
237 return i
238 res = []
239 try:
240 for x in iter(spam, 20):
241 res.append(x)
242 except RuntimeError:
Guido van Rossum805365e2007-05-07 22:24:25 +0000243 self.assertEqual(res, list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000244 else:
245 self.fail("should have raised RuntimeError")
246
247 # Test exception propagation through sequence iterator
248 def test_exception_sequence(self):
249 class MySequenceClass(SequenceClass):
250 def __getitem__(self, i):
251 if i == 10:
252 raise RuntimeError
253 return SequenceClass.__getitem__(self, i)
254 res = []
255 try:
256 for x in MySequenceClass(20):
257 res.append(x)
258 except RuntimeError:
Guido van Rossum805365e2007-05-07 22:24:25 +0000259 self.assertEqual(res, list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000260 else:
261 self.fail("should have raised RuntimeError")
262
263 # Test for StopIteration from __getitem__
264 def test_stop_sequence(self):
265 class MySequenceClass(SequenceClass):
266 def __getitem__(self, i):
267 if i == 10:
268 raise StopIteration
269 return SequenceClass.__getitem__(self, i)
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000270 self.check_for_loop(MySequenceClass(20), list(range(10)), pickle=False)
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000271
272 # Test a big range
273 def test_iter_big_range(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000274 self.check_for_loop(iter(range(10000)), list(range(10000)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000275
276 # Test an empty list
277 def test_iter_empty(self):
278 self.check_for_loop(iter([]), [])
279
280 # Test a tuple
281 def test_iter_tuple(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000282 self.check_for_loop(iter((0,1,2,3,4,5,6,7,8,9)), list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000283
Guido van Rossum805365e2007-05-07 22:24:25 +0000284 # Test a range
285 def test_iter_range(self):
286 self.check_for_loop(iter(range(10)), list(range(10)))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000287
288 # Test a string
289 def test_iter_string(self):
290 self.check_for_loop(iter("abcde"), ["a", "b", "c", "d", "e"])
291
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000292 # Test a directory
293 def test_iter_dict(self):
294 dict = {}
295 for i in range(10):
296 dict[i] = None
Brett Cannon0caf6d82007-02-22 04:49:03 +0000297 self.check_for_loop(dict, list(dict.keys()))
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000298
299 # Test a file
300 def test_iter_file(self):
301 f = open(TESTFN, "w")
302 try:
303 for i in range(5):
304 f.write("%d\n" % i)
305 finally:
306 f.close()
307 f = open(TESTFN, "r")
308 try:
Kristján Valur Jónsson31668b82012-04-03 10:49:41 +0000309 self.check_for_loop(f, ["0\n", "1\n", "2\n", "3\n", "4\n"], pickle=False)
310 self.check_for_loop(f, [], pickle=False)
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000311 finally:
312 f.close()
313 try:
314 unlink(TESTFN)
315 except OSError:
316 pass
317
Tim Petersf553f892001-05-01 20:45:31 +0000318 # Test list()'s use of iterators.
319 def test_builtin_list(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000320 self.assertEqual(list(SequenceClass(5)), list(range(5)))
Tim Petersf553f892001-05-01 20:45:31 +0000321 self.assertEqual(list(SequenceClass(0)), [])
322 self.assertEqual(list(()), [])
Tim Petersf553f892001-05-01 20:45:31 +0000323
324 d = {"one": 1, "two": 2, "three": 3}
Brett Cannon0caf6d82007-02-22 04:49:03 +0000325 self.assertEqual(list(d), list(d.keys()))
Tim Petersf553f892001-05-01 20:45:31 +0000326
327 self.assertRaises(TypeError, list, list)
328 self.assertRaises(TypeError, list, 42)
329
330 f = open(TESTFN, "w")
331 try:
332 for i in range(5):
333 f.write("%d\n" % i)
334 finally:
335 f.close()
336 f = open(TESTFN, "r")
337 try:
338 self.assertEqual(list(f), ["0\n", "1\n", "2\n", "3\n", "4\n"])
339 f.seek(0, 0)
Guido van Rossum8ee52432002-08-06 17:14:04 +0000340 self.assertEqual(list(f),
Tim Petersf553f892001-05-01 20:45:31 +0000341 ["0\n", "1\n", "2\n", "3\n", "4\n"])
342 finally:
343 f.close()
344 try:
345 unlink(TESTFN)
346 except OSError:
347 pass
348
Tim Peters6912d4d2001-05-05 03:56:37 +0000349 # Test tuples()'s use of iterators.
350 def test_builtin_tuple(self):
351 self.assertEqual(tuple(SequenceClass(5)), (0, 1, 2, 3, 4))
352 self.assertEqual(tuple(SequenceClass(0)), ())
353 self.assertEqual(tuple([]), ())
354 self.assertEqual(tuple(()), ())
355 self.assertEqual(tuple("abc"), ("a", "b", "c"))
356
357 d = {"one": 1, "two": 2, "three": 3}
358 self.assertEqual(tuple(d), tuple(d.keys()))
359
360 self.assertRaises(TypeError, tuple, list)
361 self.assertRaises(TypeError, tuple, 42)
362
363 f = open(TESTFN, "w")
364 try:
365 for i in range(5):
366 f.write("%d\n" % i)
367 finally:
368 f.close()
369 f = open(TESTFN, "r")
370 try:
371 self.assertEqual(tuple(f), ("0\n", "1\n", "2\n", "3\n", "4\n"))
372 f.seek(0, 0)
Guido van Rossum8ee52432002-08-06 17:14:04 +0000373 self.assertEqual(tuple(f),
Tim Peters6912d4d2001-05-05 03:56:37 +0000374 ("0\n", "1\n", "2\n", "3\n", "4\n"))
375 finally:
376 f.close()
377 try:
378 unlink(TESTFN)
379 except OSError:
380 pass
381
Tim Peters0e57abf2001-05-02 07:39:38 +0000382 # Test filter()'s use of iterators.
383 def test_builtin_filter(self):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000384 self.assertEqual(list(filter(None, SequenceClass(5))),
385 list(range(1, 5)))
386 self.assertEqual(list(filter(None, SequenceClass(0))), [])
387 self.assertEqual(list(filter(None, ())), [])
388 self.assertEqual(list(filter(None, "abc")), ["a", "b", "c"])
Tim Peters0e57abf2001-05-02 07:39:38 +0000389
390 d = {"one": 1, "two": 2, "three": 3}
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000391 self.assertEqual(list(filter(None, d)), list(d.keys()))
Tim Peters0e57abf2001-05-02 07:39:38 +0000392
393 self.assertRaises(TypeError, filter, None, list)
394 self.assertRaises(TypeError, filter, None, 42)
395
396 class Boolean:
397 def __init__(self, truth):
398 self.truth = truth
Jack Diederich4dafcc42006-11-28 19:15:13 +0000399 def __bool__(self):
Tim Peters0e57abf2001-05-02 07:39:38 +0000400 return self.truth
Jack Diederich4dafcc42006-11-28 19:15:13 +0000401 bTrue = Boolean(True)
402 bFalse = Boolean(False)
Tim Peters0e57abf2001-05-02 07:39:38 +0000403
404 class Seq:
405 def __init__(self, *args):
406 self.vals = args
407 def __iter__(self):
408 class SeqIter:
409 def __init__(self, vals):
410 self.vals = vals
411 self.i = 0
412 def __iter__(self):
413 return self
Georg Brandla18af4e2007-04-21 15:47:16 +0000414 def __next__(self):
Tim Peters0e57abf2001-05-02 07:39:38 +0000415 i = self.i
416 self.i = i + 1
417 if i < len(self.vals):
418 return self.vals[i]
419 else:
420 raise StopIteration
421 return SeqIter(self.vals)
422
Tim Peterscae330e2002-12-23 16:50:58 +0000423 seq = Seq(*([bTrue, bFalse] * 25))
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000424 self.assertEqual(list(filter(lambda x: not x, seq)), [bFalse]*25)
425 self.assertEqual(list(filter(lambda x: not x, iter(seq))), [bFalse]*25)
Tim Peters0e57abf2001-05-02 07:39:38 +0000426
Tim Petersc3074532001-05-03 07:00:32 +0000427 # Test max() and min()'s use of iterators.
428 def test_builtin_max_min(self):
429 self.assertEqual(max(SequenceClass(5)), 4)
430 self.assertEqual(min(SequenceClass(5)), 0)
431 self.assertEqual(max(8, -1), 8)
432 self.assertEqual(min(8, -1), -1)
433
434 d = {"one": 1, "two": 2, "three": 3}
435 self.assertEqual(max(d), "two")
436 self.assertEqual(min(d), "one")
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000437 self.assertEqual(max(d.values()), 3)
438 self.assertEqual(min(iter(d.values())), 1)
Tim Petersc3074532001-05-03 07:00:32 +0000439
Tim Petersc3074532001-05-03 07:00:32 +0000440 f = open(TESTFN, "w")
441 try:
442 f.write("medium line\n")
443 f.write("xtra large line\n")
444 f.write("itty-bitty line\n")
445 finally:
446 f.close()
447 f = open(TESTFN, "r")
448 try:
449 self.assertEqual(min(f), "itty-bitty line\n")
450 f.seek(0, 0)
451 self.assertEqual(max(f), "xtra large line\n")
452 finally:
453 f.close()
454 try:
455 unlink(TESTFN)
456 except OSError:
457 pass
458
Tim Peters4e9afdc2001-05-03 23:54:49 +0000459 # Test map()'s use of iterators.
460 def test_builtin_map(self):
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000461 self.assertEqual(list(map(lambda x: x+1, SequenceClass(5))),
462 list(range(1, 6)))
Tim Peters4e9afdc2001-05-03 23:54:49 +0000463
464 d = {"one": 1, "two": 2, "three": 3}
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000465 self.assertEqual(list(map(lambda k, d=d: (k, d[k]), d)),
466 list(d.items()))
Brett Cannon0caf6d82007-02-22 04:49:03 +0000467 dkeys = list(d.keys())
Tim Peters4e9afdc2001-05-03 23:54:49 +0000468 expected = [(i < len(d) and dkeys[i] or None,
469 i,
470 i < len(d) and dkeys[i] or None)
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000471 for i in range(3)]
Tim Peters4e9afdc2001-05-03 23:54:49 +0000472
473 f = open(TESTFN, "w")
474 try:
475 for i in range(10):
476 f.write("xy" * i + "\n") # line i has len 2*i+1
477 finally:
478 f.close()
479 f = open(TESTFN, "r")
480 try:
Guido van Rossumc1f779c2007-07-03 08:25:58 +0000481 self.assertEqual(list(map(len, f)), list(range(1, 21, 2)))
Tim Peters4e9afdc2001-05-03 23:54:49 +0000482 finally:
483 f.close()
484 try:
485 unlink(TESTFN)
486 except OSError:
487 pass
488
Tim Peters8572b4f2001-05-06 01:05:02 +0000489 # Test zip()'s use of iterators.
490 def test_builtin_zip(self):
Guido van Rossum801f0d72006-08-24 19:48:10 +0000491 self.assertEqual(list(zip()), [])
492 self.assertEqual(list(zip(*[])), [])
493 self.assertEqual(list(zip(*[(1, 2), 'ab'])), [(1, 'a'), (2, 'b')])
Raymond Hettingereaef6152003-08-02 07:42:57 +0000494
Tim Peters8572b4f2001-05-06 01:05:02 +0000495 self.assertRaises(TypeError, zip, None)
496 self.assertRaises(TypeError, zip, range(10), 42)
497 self.assertRaises(TypeError, zip, range(10), zip)
498
Guido van Rossum801f0d72006-08-24 19:48:10 +0000499 self.assertEqual(list(zip(IteratingSequenceClass(3))),
Tim Peters8572b4f2001-05-06 01:05:02 +0000500 [(0,), (1,), (2,)])
Guido van Rossum801f0d72006-08-24 19:48:10 +0000501 self.assertEqual(list(zip(SequenceClass(3))),
Tim Peters8572b4f2001-05-06 01:05:02 +0000502 [(0,), (1,), (2,)])
503
504 d = {"one": 1, "two": 2, "three": 3}
Brett Cannon0caf6d82007-02-22 04:49:03 +0000505 self.assertEqual(list(d.items()), list(zip(d, d.values())))
Tim Peters8572b4f2001-05-06 01:05:02 +0000506
507 # Generate all ints starting at constructor arg.
508 class IntsFrom:
509 def __init__(self, start):
510 self.i = start
511
512 def __iter__(self):
513 return self
514
Georg Brandla18af4e2007-04-21 15:47:16 +0000515 def __next__(self):
Tim Peters8572b4f2001-05-06 01:05:02 +0000516 i = self.i
517 self.i = i+1
518 return i
519
520 f = open(TESTFN, "w")
521 try:
522 f.write("a\n" "bbb\n" "cc\n")
523 finally:
524 f.close()
525 f = open(TESTFN, "r")
526 try:
Guido van Rossum801f0d72006-08-24 19:48:10 +0000527 self.assertEqual(list(zip(IntsFrom(0), f, IntsFrom(-100))),
Tim Peters8572b4f2001-05-06 01:05:02 +0000528 [(0, "a\n", -100),
529 (1, "bbb\n", -99),
530 (2, "cc\n", -98)])
531 finally:
532 f.close()
533 try:
534 unlink(TESTFN)
535 except OSError:
536 pass
537
Guido van Rossum805365e2007-05-07 22:24:25 +0000538 self.assertEqual(list(zip(range(5))), [(i,) for i in range(5)])
Tim Peters67d687a2002-04-29 21:27:32 +0000539
540 # Classes that lie about their lengths.
541 class NoGuessLen5:
542 def __getitem__(self, i):
543 if i >= 5:
544 raise IndexError
545 return i
546
547 class Guess3Len5(NoGuessLen5):
548 def __len__(self):
549 return 3
550
551 class Guess30Len5(NoGuessLen5):
552 def __len__(self):
553 return 30
554
Guido van Rossum801f0d72006-08-24 19:48:10 +0000555 def lzip(*args):
556 return list(zip(*args))
557
Tim Peters67d687a2002-04-29 21:27:32 +0000558 self.assertEqual(len(Guess3Len5()), 3)
559 self.assertEqual(len(Guess30Len5()), 30)
Guido van Rossum801f0d72006-08-24 19:48:10 +0000560 self.assertEqual(lzip(NoGuessLen5()), lzip(range(5)))
561 self.assertEqual(lzip(Guess3Len5()), lzip(range(5)))
562 self.assertEqual(lzip(Guess30Len5()), lzip(range(5)))
Tim Peters67d687a2002-04-29 21:27:32 +0000563
564 expected = [(i, i) for i in range(5)]
565 for x in NoGuessLen5(), Guess3Len5(), Guess30Len5():
566 for y in NoGuessLen5(), Guess3Len5(), Guess30Len5():
Guido van Rossum801f0d72006-08-24 19:48:10 +0000567 self.assertEqual(lzip(x, y), expected)
Tim Peters67d687a2002-04-29 21:27:32 +0000568
Tim Peters2cfe3682001-05-05 05:36:48 +0000569 def test_unicode_join_endcase(self):
570
571 # This class inserts a Unicode object into its argument's natural
572 # iteration, in the 3rd position.
573 class OhPhooey:
574 def __init__(self, seq):
575 self.it = iter(seq)
576 self.i = 0
577
578 def __iter__(self):
579 return self
580
Georg Brandla18af4e2007-04-21 15:47:16 +0000581 def __next__(self):
Tim Peters2cfe3682001-05-05 05:36:48 +0000582 i = self.i
583 self.i = i+1
584 if i == 2:
Walter Dörwald1f5947b2007-05-22 16:52:54 +0000585 return "fooled you!"
Georg Brandla18af4e2007-04-21 15:47:16 +0000586 return next(self.it)
Tim Peters2cfe3682001-05-05 05:36:48 +0000587
588 f = open(TESTFN, "w")
589 try:
590 f.write("a\n" + "b\n" + "c\n")
591 finally:
592 f.close()
593
594 f = open(TESTFN, "r")
595 # Nasty: string.join(s) can't know whether unicode.join() is needed
596 # until it's seen all of s's elements. But in this case, f's
597 # iterator cannot be restarted. So what we're testing here is
598 # whether string.join() can manage to remember everything it's seen
599 # and pass that on to unicode.join().
600 try:
601 got = " - ".join(OhPhooey(f))
Walter Dörwald5de48bd2007-06-11 21:38:39 +0000602 self.assertEqual(got, "a\n - b\n - fooled you! - c\n")
Tim Peters2cfe3682001-05-05 05:36:48 +0000603 finally:
604 f.close()
605 try:
606 unlink(TESTFN)
607 except OSError:
608 pass
609
Tim Petersde9725f2001-05-05 10:06:17 +0000610 # Test iterators with 'x in y' and 'x not in y'.
611 def test_in_and_not_in(self):
Tim Peterscb8d3682001-05-05 21:05:01 +0000612 for sc5 in IteratingSequenceClass(5), SequenceClass(5):
613 for i in range(5):
Benjamin Peterson577473f2010-01-19 00:09:57 +0000614 self.assertIn(i, sc5)
Tim Peterscb8d3682001-05-05 21:05:01 +0000615 for i in "abc", -1, 5, 42.42, (3, 4), [], {1: 1}, 3-12j, sc5:
Benjamin Peterson577473f2010-01-19 00:09:57 +0000616 self.assertNotIn(i, sc5)
Tim Petersde9725f2001-05-05 10:06:17 +0000617
618 self.assertRaises(TypeError, lambda: 3 in 12)
619 self.assertRaises(TypeError, lambda: 3 not in map)
620
621 d = {"one": 1, "two": 2, "three": 3, 1j: 2j}
622 for k in d:
Benjamin Peterson577473f2010-01-19 00:09:57 +0000623 self.assertIn(k, d)
624 self.assertNotIn(k, d.values())
Tim Petersde9725f2001-05-05 10:06:17 +0000625 for v in d.values():
Benjamin Peterson577473f2010-01-19 00:09:57 +0000626 self.assertIn(v, d.values())
627 self.assertNotIn(v, d)
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000628 for k, v in d.items():
Benjamin Peterson577473f2010-01-19 00:09:57 +0000629 self.assertIn((k, v), d.items())
630 self.assertNotIn((v, k), d.items())
Tim Petersde9725f2001-05-05 10:06:17 +0000631
632 f = open(TESTFN, "w")
633 try:
634 f.write("a\n" "b\n" "c\n")
635 finally:
636 f.close()
637 f = open(TESTFN, "r")
638 try:
639 for chunk in "abc":
640 f.seek(0, 0)
Benjamin Peterson577473f2010-01-19 00:09:57 +0000641 self.assertNotIn(chunk, f)
Tim Petersde9725f2001-05-05 10:06:17 +0000642 f.seek(0, 0)
Benjamin Peterson577473f2010-01-19 00:09:57 +0000643 self.assertIn((chunk + "\n"), f)
Tim Petersde9725f2001-05-05 10:06:17 +0000644 finally:
645 f.close()
646 try:
647 unlink(TESTFN)
648 except OSError:
649 pass
650
Tim Peters75f8e352001-05-05 11:33:43 +0000651 # Test iterators with operator.countOf (PySequence_Count).
652 def test_countOf(self):
653 from operator import countOf
654 self.assertEqual(countOf([1,2,2,3,2,5], 2), 3)
655 self.assertEqual(countOf((1,2,2,3,2,5), 2), 3)
656 self.assertEqual(countOf("122325", "2"), 3)
657 self.assertEqual(countOf("122325", "6"), 0)
658
659 self.assertRaises(TypeError, countOf, 42, 1)
660 self.assertRaises(TypeError, countOf, countOf, countOf)
661
662 d = {"one": 3, "two": 3, "three": 3, 1j: 2j}
663 for k in d:
664 self.assertEqual(countOf(d, k), 1)
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000665 self.assertEqual(countOf(d.values(), 3), 3)
666 self.assertEqual(countOf(d.values(), 2j), 1)
667 self.assertEqual(countOf(d.values(), 1j), 0)
Tim Peters75f8e352001-05-05 11:33:43 +0000668
669 f = open(TESTFN, "w")
670 try:
671 f.write("a\n" "b\n" "c\n" "b\n")
672 finally:
673 f.close()
674 f = open(TESTFN, "r")
675 try:
676 for letter, count in ("a", 1), ("b", 2), ("c", 1), ("d", 0):
677 f.seek(0, 0)
678 self.assertEqual(countOf(f, letter + "\n"), count)
679 finally:
680 f.close()
681 try:
682 unlink(TESTFN)
683 except OSError:
684 pass
685
Tim Peters16a77ad2001-09-08 04:00:12 +0000686 # Test iterators with operator.indexOf (PySequence_Index).
687 def test_indexOf(self):
688 from operator import indexOf
689 self.assertEqual(indexOf([1,2,2,3,2,5], 1), 0)
690 self.assertEqual(indexOf((1,2,2,3,2,5), 2), 1)
691 self.assertEqual(indexOf((1,2,2,3,2,5), 3), 3)
692 self.assertEqual(indexOf((1,2,2,3,2,5), 5), 5)
693 self.assertRaises(ValueError, indexOf, (1,2,2,3,2,5), 0)
694 self.assertRaises(ValueError, indexOf, (1,2,2,3,2,5), 6)
695
696 self.assertEqual(indexOf("122325", "2"), 1)
697 self.assertEqual(indexOf("122325", "5"), 5)
698 self.assertRaises(ValueError, indexOf, "122325", "6")
699
700 self.assertRaises(TypeError, indexOf, 42, 1)
701 self.assertRaises(TypeError, indexOf, indexOf, indexOf)
702
703 f = open(TESTFN, "w")
704 try:
705 f.write("a\n" "b\n" "c\n" "d\n" "e\n")
706 finally:
707 f.close()
708 f = open(TESTFN, "r")
709 try:
710 fiter = iter(f)
711 self.assertEqual(indexOf(fiter, "b\n"), 1)
712 self.assertEqual(indexOf(fiter, "d\n"), 1)
713 self.assertEqual(indexOf(fiter, "e\n"), 0)
714 self.assertRaises(ValueError, indexOf, fiter, "a\n")
715 finally:
716 f.close()
717 try:
718 unlink(TESTFN)
719 except OSError:
720 pass
721
722 iclass = IteratingSequenceClass(3)
723 for i in range(3):
724 self.assertEqual(indexOf(iclass, i), i)
725 self.assertRaises(ValueError, indexOf, iclass, -1)
726
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000727 # Test iterators with file.writelines().
728 def test_writelines(self):
Alex Martelli01c77c62006-08-24 02:58:11 +0000729 f = open(TESTFN, "w")
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000730
731 try:
732 self.assertRaises(TypeError, f.writelines, None)
733 self.assertRaises(TypeError, f.writelines, 42)
Tim Peters527e64f2001-10-04 05:36:56 +0000734
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000735 f.writelines(["1\n", "2\n"])
736 f.writelines(("3\n", "4\n"))
737 f.writelines({'5\n': None})
738 f.writelines({})
739
740 # Try a big chunk too.
741 class Iterator:
742 def __init__(self, start, finish):
743 self.start = start
744 self.finish = finish
745 self.i = self.start
746
Georg Brandla18af4e2007-04-21 15:47:16 +0000747 def __next__(self):
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000748 if self.i >= self.finish:
749 raise StopIteration
750 result = str(self.i) + '\n'
751 self.i += 1
752 return result
753
754 def __iter__(self):
755 return self
756
757 class Whatever:
758 def __init__(self, start, finish):
759 self.start = start
760 self.finish = finish
761
762 def __iter__(self):
763 return Iterator(self.start, self.finish)
Tim Peters527e64f2001-10-04 05:36:56 +0000764
765 f.writelines(Whatever(6, 6+2000))
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000766 f.close()
767
Alex Martelli01c77c62006-08-24 02:58:11 +0000768 f = open(TESTFN)
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000769 expected = [str(i) + "\n" for i in range(1, 2006)]
770 self.assertEqual(list(f), expected)
Tim Peters527e64f2001-10-04 05:36:56 +0000771
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000772 finally:
773 f.close()
774 try:
775 unlink(TESTFN)
776 except OSError:
777 pass
778
779
Tim Petersd6d010b2001-06-21 02:49:55 +0000780 # Test iterators on RHS of unpacking assignments.
781 def test_unpack_iter(self):
782 a, b = 1, 2
783 self.assertEqual((a, b), (1, 2))
784
785 a, b, c = IteratingSequenceClass(3)
786 self.assertEqual((a, b, c), (0, 1, 2))
787
788 try: # too many values
789 a, b = IteratingSequenceClass(3)
790 except ValueError:
791 pass
792 else:
793 self.fail("should have raised ValueError")
794
795 try: # not enough values
796 a, b, c = IteratingSequenceClass(2)
797 except ValueError:
798 pass
799 else:
800 self.fail("should have raised ValueError")
801
802 try: # not iterable
803 a, b, c = len
804 except TypeError:
805 pass
806 else:
807 self.fail("should have raised TypeError")
808
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000809 a, b, c = {1: 42, 2: 42, 3: 42}.values()
Tim Petersd6d010b2001-06-21 02:49:55 +0000810 self.assertEqual((a, b, c), (42, 42, 42))
811
812 f = open(TESTFN, "w")
813 lines = ("a\n", "bb\n", "ccc\n")
814 try:
815 for line in lines:
816 f.write(line)
817 finally:
818 f.close()
819 f = open(TESTFN, "r")
820 try:
821 a, b, c = f
822 self.assertEqual((a, b, c), lines)
823 finally:
824 f.close()
825 try:
826 unlink(TESTFN)
827 except OSError:
828 pass
829
830 (a, b), (c,) = IteratingSequenceClass(2), {42: 24}
831 self.assertEqual((a, b, c), (0, 1, 42))
832
Guido van Rossumbb8f59a2001-12-03 19:33:25 +0000833
Benjamin Peterson945c5792010-06-22 20:34:34 +0000834 @cpython_only
835 def test_ref_counting_behavior(self):
Guido van Rossumbb8f59a2001-12-03 19:33:25 +0000836 class C(object):
837 count = 0
838 def __new__(cls):
839 cls.count += 1
840 return object.__new__(cls)
841 def __del__(self):
842 cls = self.__class__
843 assert cls.count > 0
844 cls.count -= 1
845 x = C()
846 self.assertEqual(C.count, 1)
847 del x
848 self.assertEqual(C.count, 0)
849 l = [C(), C(), C()]
850 self.assertEqual(C.count, 3)
851 try:
852 a, b = iter(l)
853 except ValueError:
854 pass
855 del l
856 self.assertEqual(C.count, 0)
Fred Drake2e2be372001-09-20 21:33:42 +0000857
Guido van Rossum674eae62002-07-16 21:48:11 +0000858
859 # Make sure StopIteration is a "sink state".
860 # This tests various things that weren't sink states in Python 2.2.1,
861 # plus various things that always were fine.
862
863 def test_sinkstate_list(self):
864 # This used to fail
Guido van Rossum805365e2007-05-07 22:24:25 +0000865 a = list(range(5))
Guido van Rossum674eae62002-07-16 21:48:11 +0000866 b = iter(a)
Guido van Rossum805365e2007-05-07 22:24:25 +0000867 self.assertEqual(list(b), list(range(5)))
Guido van Rossum674eae62002-07-16 21:48:11 +0000868 a.extend(range(5, 10))
869 self.assertEqual(list(b), [])
870
871 def test_sinkstate_tuple(self):
872 a = (0, 1, 2, 3, 4)
873 b = iter(a)
Guido van Rossum805365e2007-05-07 22:24:25 +0000874 self.assertEqual(list(b), list(range(5)))
Guido van Rossum674eae62002-07-16 21:48:11 +0000875 self.assertEqual(list(b), [])
876
877 def test_sinkstate_string(self):
878 a = "abcde"
879 b = iter(a)
880 self.assertEqual(list(b), ['a', 'b', 'c', 'd', 'e'])
881 self.assertEqual(list(b), [])
882
883 def test_sinkstate_sequence(self):
884 # This used to fail
885 a = SequenceClass(5)
886 b = iter(a)
Guido van Rossum805365e2007-05-07 22:24:25 +0000887 self.assertEqual(list(b), list(range(5)))
Guido van Rossum674eae62002-07-16 21:48:11 +0000888 a.n = 10
889 self.assertEqual(list(b), [])
890
891 def test_sinkstate_callable(self):
892 # This used to fail
893 def spam(state=[0]):
894 i = state[0]
895 state[0] = i+1
896 if i == 10:
Collin Winter3add4d72007-08-29 23:37:32 +0000897 raise AssertionError("shouldn't have gotten this far")
Guido van Rossum674eae62002-07-16 21:48:11 +0000898 return i
899 b = iter(spam, 5)
Guido van Rossum805365e2007-05-07 22:24:25 +0000900 self.assertEqual(list(b), list(range(5)))
Guido van Rossum674eae62002-07-16 21:48:11 +0000901 self.assertEqual(list(b), [])
902
903 def test_sinkstate_dict(self):
904 # XXX For a more thorough test, see towards the end of:
905 # http://mail.python.org/pipermail/python-dev/2002-July/026512.html
906 a = {1:1, 2:2, 0:0, 4:4, 3:3}
Guido van Rossumcc2b0162007-02-11 06:12:03 +0000907 for b in iter(a), a.keys(), a.items(), a.values():
Guido van Rossum674eae62002-07-16 21:48:11 +0000908 b = iter(a)
909 self.assertEqual(len(list(b)), 5)
910 self.assertEqual(list(b), [])
911
912 def test_sinkstate_yield(self):
913 def gen():
914 for i in range(5):
915 yield i
916 b = gen()
Guido van Rossum805365e2007-05-07 22:24:25 +0000917 self.assertEqual(list(b), list(range(5)))
Guido van Rossum674eae62002-07-16 21:48:11 +0000918 self.assertEqual(list(b), [])
919
920 def test_sinkstate_range(self):
Guido van Rossum805365e2007-05-07 22:24:25 +0000921 a = range(5)
Guido van Rossum674eae62002-07-16 21:48:11 +0000922 b = iter(a)
Guido van Rossum805365e2007-05-07 22:24:25 +0000923 self.assertEqual(list(b), list(range(5)))
Guido van Rossum674eae62002-07-16 21:48:11 +0000924 self.assertEqual(list(b), [])
925
926 def test_sinkstate_enumerate(self):
927 a = range(5)
928 e = enumerate(a)
929 b = iter(e)
Guido van Rossum801f0d72006-08-24 19:48:10 +0000930 self.assertEqual(list(b), list(zip(range(5), range(5))))
Guido van Rossum674eae62002-07-16 21:48:11 +0000931 self.assertEqual(list(b), [])
932
Amaury Forgeot d'Arcf343e012009-01-12 23:58:21 +0000933 def test_3720(self):
934 # Avoid a crash, when an iterator deletes its next() method.
935 class BadIterator(object):
936 def __iter__(self):
937 return self
938 def __next__(self):
939 del BadIterator.__next__
940 return 1
941
942 try:
943 for i in BadIterator() :
944 pass
945 except TypeError:
946 pass
947
Ezio Melotti739e1792012-11-18 23:16:02 +0200948 def test_extending_list_with_iterator_does_not_segfault(self):
949 # The code to extend a list with an iterator has a fair
950 # amount of nontrivial logic in terms of guessing how
951 # much memory to allocate in advance, "stealing" refs,
952 # and then shrinking at the end. This is a basic smoke
953 # test for that scenario.
954 def gen():
955 for i in range(500):
956 yield i
957 lst = [0] * 500
958 for i in range(240):
959 lst.pop(0)
960 lst.extend(gen())
961 self.assertEqual(len(lst), 760)
962
Serhiy Storchaka4faf5c52015-05-21 20:50:25 +0300963 @cpython_only
964 def test_iter_overflow(self):
965 # Test for the issue 22939
966 it = iter(UnlimitedSequenceClass())
967 # Manually set `it_index` to PY_SSIZE_T_MAX-2 without a loop
968 it.__setstate__(sys.maxsize - 2)
969 self.assertEqual(next(it), sys.maxsize - 2)
970 self.assertEqual(next(it), sys.maxsize - 1)
971 with self.assertRaises(OverflowError):
972 next(it)
973 # Check that Overflow error is always raised
974 with self.assertRaises(OverflowError):
975 next(it)
976
977 def test_iter_neg_setstate(self):
978 it = iter(UnlimitedSequenceClass())
979 it.__setstate__(-42)
980 self.assertEqual(next(it), 0)
981 self.assertEqual(next(it), 1)
982
Guido van Rossum674eae62002-07-16 21:48:11 +0000983
Fred Drake2e2be372001-09-20 21:33:42 +0000984def test_main():
985 run_unittest(TestCase)
986
987
988if __name__ == "__main__":
989 test_main()