blob: 0085db5dc129ea307b10929e3d142426199c961b [file] [log] [blame]
Guido van Rossum8b48cf92001-04-21 13:33:54 +00001# Test iterators.
2
3import unittest
Martin v. Löwis339d0f72001-08-17 18:39:25 +00004from test_support import run_unittest, TESTFN, unlink, have_unicode
Guido van Rossum8b48cf92001-04-21 13:33:54 +00005
6# Test result of triple loop (too big to inline)
7TRIPLETS = [(0, 0, 0), (0, 0, 1), (0, 0, 2),
8 (0, 1, 0), (0, 1, 1), (0, 1, 2),
9 (0, 2, 0), (0, 2, 1), (0, 2, 2),
10
11 (1, 0, 0), (1, 0, 1), (1, 0, 2),
12 (1, 1, 0), (1, 1, 1), (1, 1, 2),
13 (1, 2, 0), (1, 2, 1), (1, 2, 2),
14
15 (2, 0, 0), (2, 0, 1), (2, 0, 2),
16 (2, 1, 0), (2, 1, 1), (2, 1, 2),
17 (2, 2, 0), (2, 2, 1), (2, 2, 2)]
18
19# Helper classes
20
21class BasicIterClass:
22 def __init__(self, n):
23 self.n = n
24 self.i = 0
25 def next(self):
26 res = self.i
27 if res >= self.n:
28 raise StopIteration
29 self.i = res + 1
30 return res
31
32class IteratingSequenceClass:
33 def __init__(self, n):
34 self.n = n
35 def __iter__(self):
36 return BasicIterClass(self.n)
37
38class SequenceClass:
39 def __init__(self, n):
40 self.n = n
41 def __getitem__(self, i):
42 if 0 <= i < self.n:
43 return i
44 else:
45 raise IndexError
46
47# Main test suite
48
49class TestCase(unittest.TestCase):
50
51 # Helper to check that an iterator returns a given sequence
52 def check_iterator(self, it, seq):
53 res = []
54 while 1:
55 try:
56 val = it.next()
57 except StopIteration:
58 break
59 res.append(val)
60 self.assertEqual(res, seq)
61
62 # Helper to check that a for loop generates a given sequence
63 def check_for_loop(self, expr, seq):
64 res = []
65 for val in expr:
66 res.append(val)
67 self.assertEqual(res, seq)
68
69 # Test basic use of iter() function
70 def test_iter_basic(self):
71 self.check_iterator(iter(range(10)), range(10))
72
73 # Test that iter(iter(x)) is the same as iter(x)
74 def test_iter_idempotency(self):
75 seq = range(10)
76 it = iter(seq)
77 it2 = iter(it)
78 self.assert_(it is it2)
79
80 # Test that for loops over iterators work
81 def test_iter_for_loop(self):
82 self.check_for_loop(iter(range(10)), range(10))
83
84 # Test several independent iterators over the same list
85 def test_iter_independence(self):
86 seq = range(3)
87 res = []
88 for i in iter(seq):
89 for j in iter(seq):
90 for k in iter(seq):
91 res.append((i, j, k))
92 self.assertEqual(res, TRIPLETS)
93
94 # Test triple list comprehension using iterators
95 def test_nested_comprehensions_iter(self):
96 seq = range(3)
97 res = [(i, j, k)
98 for i in iter(seq) for j in iter(seq) for k in iter(seq)]
99 self.assertEqual(res, TRIPLETS)
100
101 # Test triple list comprehension without iterators
102 def test_nested_comprehensions_for(self):
103 seq = range(3)
104 res = [(i, j, k) for i in seq for j in seq for k in seq]
105 self.assertEqual(res, TRIPLETS)
106
107 # Test a class with __iter__ in a for loop
108 def test_iter_class_for(self):
109 self.check_for_loop(IteratingSequenceClass(10), range(10))
110
111 # Test a class with __iter__ with explicit iter()
112 def test_iter_class_iter(self):
113 self.check_iterator(iter(IteratingSequenceClass(10)), range(10))
114
115 # Test for loop on a sequence class without __iter__
116 def test_seq_class_for(self):
117 self.check_for_loop(SequenceClass(10), range(10))
118
119 # Test iter() on a sequence class without __iter__
120 def test_seq_class_iter(self):
121 self.check_iterator(iter(SequenceClass(10)), range(10))
122
123 # Test two-argument iter() with callable instance
124 def test_iter_callable(self):
125 class C:
126 def __init__(self):
127 self.i = 0
128 def __call__(self):
129 i = self.i
130 self.i = i + 1
131 if i > 100:
132 raise IndexError # Emergency stop
133 return i
134 self.check_iterator(iter(C(), 10), range(10))
135
136 # Test two-argument iter() with function
137 def test_iter_function(self):
138 def spam(state=[0]):
139 i = state[0]
140 state[0] = i+1
141 return i
142 self.check_iterator(iter(spam, 10), range(10))
143
144 # Test two-argument iter() with function that raises StopIteration
145 def test_iter_function_stop(self):
146 def spam(state=[0]):
147 i = state[0]
148 if i == 10:
149 raise StopIteration
150 state[0] = i+1
151 return i
152 self.check_iterator(iter(spam, 20), range(10))
153
154 # Test exception propagation through function iterator
155 def test_exception_function(self):
156 def spam(state=[0]):
157 i = state[0]
158 state[0] = i+1
159 if i == 10:
160 raise RuntimeError
161 return i
162 res = []
163 try:
164 for x in iter(spam, 20):
165 res.append(x)
166 except RuntimeError:
167 self.assertEqual(res, range(10))
168 else:
169 self.fail("should have raised RuntimeError")
170
171 # Test exception propagation through sequence iterator
172 def test_exception_sequence(self):
173 class MySequenceClass(SequenceClass):
174 def __getitem__(self, i):
175 if i == 10:
176 raise RuntimeError
177 return SequenceClass.__getitem__(self, i)
178 res = []
179 try:
180 for x in MySequenceClass(20):
181 res.append(x)
182 except RuntimeError:
183 self.assertEqual(res, range(10))
184 else:
185 self.fail("should have raised RuntimeError")
186
187 # Test for StopIteration from __getitem__
188 def test_stop_sequence(self):
189 class MySequenceClass(SequenceClass):
190 def __getitem__(self, i):
191 if i == 10:
192 raise StopIteration
193 return SequenceClass.__getitem__(self, i)
194 self.check_for_loop(MySequenceClass(20), range(10))
195
196 # Test a big range
197 def test_iter_big_range(self):
198 self.check_for_loop(iter(range(10000)), range(10000))
199
200 # Test an empty list
201 def test_iter_empty(self):
202 self.check_for_loop(iter([]), [])
203
204 # Test a tuple
205 def test_iter_tuple(self):
206 self.check_for_loop(iter((0,1,2,3,4,5,6,7,8,9)), range(10))
207
208 # Test an xrange
209 def test_iter_xrange(self):
210 self.check_for_loop(iter(xrange(10)), range(10))
211
212 # Test a string
213 def test_iter_string(self):
214 self.check_for_loop(iter("abcde"), ["a", "b", "c", "d", "e"])
215
216 # Test a Unicode string
Martin v. Löwis339d0f72001-08-17 18:39:25 +0000217 if have_unicode:
218 def test_iter_unicode(self):
219 self.check_for_loop(iter(unicode("abcde")),
220 [unicode("a"), unicode("b"), unicode("c"),
221 unicode("d"), unicode("e")])
Guido van Rossum8b48cf92001-04-21 13:33:54 +0000222
223 # Test a directory
224 def test_iter_dict(self):
225 dict = {}
226 for i in range(10):
227 dict[i] = None
228 self.check_for_loop(dict, dict.keys())
229
230 # Test a file
231 def test_iter_file(self):
232 f = open(TESTFN, "w")
233 try:
234 for i in range(5):
235 f.write("%d\n" % i)
236 finally:
237 f.close()
238 f = open(TESTFN, "r")
239 try:
240 self.check_for_loop(f, ["0\n", "1\n", "2\n", "3\n", "4\n"])
241 self.check_for_loop(f, [])
242 finally:
243 f.close()
244 try:
245 unlink(TESTFN)
246 except OSError:
247 pass
248
Tim Petersf553f892001-05-01 20:45:31 +0000249 # Test list()'s use of iterators.
250 def test_builtin_list(self):
251 self.assertEqual(list(SequenceClass(5)), range(5))
252 self.assertEqual(list(SequenceClass(0)), [])
253 self.assertEqual(list(()), [])
254 self.assertEqual(list(range(10, -1, -1)), range(10, -1, -1))
255
256 d = {"one": 1, "two": 2, "three": 3}
257 self.assertEqual(list(d), d.keys())
258
259 self.assertRaises(TypeError, list, list)
260 self.assertRaises(TypeError, list, 42)
261
262 f = open(TESTFN, "w")
263 try:
264 for i in range(5):
265 f.write("%d\n" % i)
266 finally:
267 f.close()
268 f = open(TESTFN, "r")
269 try:
270 self.assertEqual(list(f), ["0\n", "1\n", "2\n", "3\n", "4\n"])
271 f.seek(0, 0)
272 self.assertEqual(list(f.xreadlines()),
273 ["0\n", "1\n", "2\n", "3\n", "4\n"])
274 finally:
275 f.close()
276 try:
277 unlink(TESTFN)
278 except OSError:
279 pass
280
Tim Peters6912d4d2001-05-05 03:56:37 +0000281 # Test tuples()'s use of iterators.
282 def test_builtin_tuple(self):
283 self.assertEqual(tuple(SequenceClass(5)), (0, 1, 2, 3, 4))
284 self.assertEqual(tuple(SequenceClass(0)), ())
285 self.assertEqual(tuple([]), ())
286 self.assertEqual(tuple(()), ())
287 self.assertEqual(tuple("abc"), ("a", "b", "c"))
288
289 d = {"one": 1, "two": 2, "three": 3}
290 self.assertEqual(tuple(d), tuple(d.keys()))
291
292 self.assertRaises(TypeError, tuple, list)
293 self.assertRaises(TypeError, tuple, 42)
294
295 f = open(TESTFN, "w")
296 try:
297 for i in range(5):
298 f.write("%d\n" % i)
299 finally:
300 f.close()
301 f = open(TESTFN, "r")
302 try:
303 self.assertEqual(tuple(f), ("0\n", "1\n", "2\n", "3\n", "4\n"))
304 f.seek(0, 0)
305 self.assertEqual(tuple(f.xreadlines()),
306 ("0\n", "1\n", "2\n", "3\n", "4\n"))
307 finally:
308 f.close()
309 try:
310 unlink(TESTFN)
311 except OSError:
312 pass
313
Tim Peters0e57abf2001-05-02 07:39:38 +0000314 # Test filter()'s use of iterators.
315 def test_builtin_filter(self):
316 self.assertEqual(filter(None, SequenceClass(5)), range(1, 5))
317 self.assertEqual(filter(None, SequenceClass(0)), [])
318 self.assertEqual(filter(None, ()), ())
319 self.assertEqual(filter(None, "abc"), "abc")
320
321 d = {"one": 1, "two": 2, "three": 3}
322 self.assertEqual(filter(None, d), d.keys())
323
324 self.assertRaises(TypeError, filter, None, list)
325 self.assertRaises(TypeError, filter, None, 42)
326
327 class Boolean:
328 def __init__(self, truth):
329 self.truth = truth
330 def __nonzero__(self):
331 return self.truth
332 True = Boolean(1)
333 False = Boolean(0)
334
335 class Seq:
336 def __init__(self, *args):
337 self.vals = args
338 def __iter__(self):
339 class SeqIter:
340 def __init__(self, vals):
341 self.vals = vals
342 self.i = 0
343 def __iter__(self):
344 return self
345 def next(self):
346 i = self.i
347 self.i = i + 1
348 if i < len(self.vals):
349 return self.vals[i]
350 else:
351 raise StopIteration
352 return SeqIter(self.vals)
353
354 seq = Seq(*([True, False] * 25))
355 self.assertEqual(filter(lambda x: not x, seq), [False]*25)
356 self.assertEqual(filter(lambda x: not x, iter(seq)), [False]*25)
357
Tim Petersc3074532001-05-03 07:00:32 +0000358 # Test max() and min()'s use of iterators.
359 def test_builtin_max_min(self):
360 self.assertEqual(max(SequenceClass(5)), 4)
361 self.assertEqual(min(SequenceClass(5)), 0)
362 self.assertEqual(max(8, -1), 8)
363 self.assertEqual(min(8, -1), -1)
364
365 d = {"one": 1, "two": 2, "three": 3}
366 self.assertEqual(max(d), "two")
367 self.assertEqual(min(d), "one")
368 self.assertEqual(max(d.itervalues()), 3)
369 self.assertEqual(min(iter(d.itervalues())), 1)
370
Tim Petersc3074532001-05-03 07:00:32 +0000371 f = open(TESTFN, "w")
372 try:
373 f.write("medium line\n")
374 f.write("xtra large line\n")
375 f.write("itty-bitty line\n")
376 finally:
377 f.close()
378 f = open(TESTFN, "r")
379 try:
380 self.assertEqual(min(f), "itty-bitty line\n")
381 f.seek(0, 0)
382 self.assertEqual(max(f), "xtra large line\n")
383 finally:
384 f.close()
385 try:
386 unlink(TESTFN)
387 except OSError:
388 pass
389
Tim Peters4e9afdc2001-05-03 23:54:49 +0000390 # Test map()'s use of iterators.
391 def test_builtin_map(self):
392 self.assertEqual(map(None, SequenceClass(5)), range(5))
393 self.assertEqual(map(lambda x: x+1, SequenceClass(5)), range(1, 6))
394
395 d = {"one": 1, "two": 2, "three": 3}
396 self.assertEqual(map(None, d), d.keys())
397 self.assertEqual(map(lambda k, d=d: (k, d[k]), d), d.items())
398 dkeys = d.keys()
399 expected = [(i < len(d) and dkeys[i] or None,
400 i,
401 i < len(d) and dkeys[i] or None)
402 for i in range(5)]
403 self.assertEqual(map(None, d,
404 SequenceClass(5),
405 iter(d.iterkeys())),
Tim Peters8bc10b02001-05-03 23:58:47 +0000406 expected)
Tim Peters4e9afdc2001-05-03 23:54:49 +0000407
408 f = open(TESTFN, "w")
409 try:
410 for i in range(10):
411 f.write("xy" * i + "\n") # line i has len 2*i+1
412 finally:
413 f.close()
414 f = open(TESTFN, "r")
415 try:
416 self.assertEqual(map(len, f), range(1, 21, 2))
Tim Peters4e9afdc2001-05-03 23:54:49 +0000417 finally:
418 f.close()
419 try:
420 unlink(TESTFN)
421 except OSError:
422 pass
423
Tim Peters8572b4f2001-05-06 01:05:02 +0000424 # Test zip()'s use of iterators.
425 def test_builtin_zip(self):
426 self.assertRaises(TypeError, zip)
427 self.assertRaises(TypeError, zip, None)
428 self.assertRaises(TypeError, zip, range(10), 42)
429 self.assertRaises(TypeError, zip, range(10), zip)
430
431 self.assertEqual(zip(IteratingSequenceClass(3)),
432 [(0,), (1,), (2,)])
433 self.assertEqual(zip(SequenceClass(3)),
434 [(0,), (1,), (2,)])
435
436 d = {"one": 1, "two": 2, "three": 3}
437 self.assertEqual(d.items(), zip(d, d.itervalues()))
438
439 # Generate all ints starting at constructor arg.
440 class IntsFrom:
441 def __init__(self, start):
442 self.i = start
443
444 def __iter__(self):
445 return self
446
447 def next(self):
448 i = self.i
449 self.i = i+1
450 return i
451
452 f = open(TESTFN, "w")
453 try:
454 f.write("a\n" "bbb\n" "cc\n")
455 finally:
456 f.close()
457 f = open(TESTFN, "r")
458 try:
459 self.assertEqual(zip(IntsFrom(0), f, IntsFrom(-100)),
460 [(0, "a\n", -100),
461 (1, "bbb\n", -99),
462 (2, "cc\n", -98)])
463 finally:
464 f.close()
465 try:
466 unlink(TESTFN)
467 except OSError:
468 pass
469
Tim Peters67d687a2002-04-29 21:27:32 +0000470 self.assertEqual(zip(xrange(5)), [(i,) for i in range(5)])
471
472 # Classes that lie about their lengths.
473 class NoGuessLen5:
474 def __getitem__(self, i):
475 if i >= 5:
476 raise IndexError
477 return i
478
479 class Guess3Len5(NoGuessLen5):
480 def __len__(self):
481 return 3
482
483 class Guess30Len5(NoGuessLen5):
484 def __len__(self):
485 return 30
486
487 self.assertEqual(len(Guess3Len5()), 3)
488 self.assertEqual(len(Guess30Len5()), 30)
489 self.assertEqual(zip(NoGuessLen5()), zip(range(5)))
490 self.assertEqual(zip(Guess3Len5()), zip(range(5)))
491 self.assertEqual(zip(Guess30Len5()), zip(range(5)))
492
493 expected = [(i, i) for i in range(5)]
494 for x in NoGuessLen5(), Guess3Len5(), Guess30Len5():
495 for y in NoGuessLen5(), Guess3Len5(), Guess30Len5():
496 self.assertEqual(zip(x, y), expected)
497
Tim Peters15d81ef2001-05-04 04:39:21 +0000498 # Test reduces()'s use of iterators.
499 def test_builtin_reduce(self):
500 from operator import add
501 self.assertEqual(reduce(add, SequenceClass(5)), 10)
502 self.assertEqual(reduce(add, SequenceClass(5), 42), 52)
503 self.assertRaises(TypeError, reduce, add, SequenceClass(0))
504 self.assertEqual(reduce(add, SequenceClass(0), 42), 42)
505 self.assertEqual(reduce(add, SequenceClass(1)), 0)
506 self.assertEqual(reduce(add, SequenceClass(1), 42), 42)
507
508 d = {"one": 1, "two": 2, "three": 3}
509 self.assertEqual(reduce(add, d), "".join(d.keys()))
510
Martin v. Löwis339d0f72001-08-17 18:39:25 +0000511 # This test case will be removed if we don't have Unicode
Tim Peters2cfe3682001-05-05 05:36:48 +0000512 def test_unicode_join_endcase(self):
513
514 # This class inserts a Unicode object into its argument's natural
515 # iteration, in the 3rd position.
516 class OhPhooey:
517 def __init__(self, seq):
518 self.it = iter(seq)
519 self.i = 0
520
521 def __iter__(self):
522 return self
523
524 def next(self):
525 i = self.i
526 self.i = i+1
527 if i == 2:
Martin v. Löwis339d0f72001-08-17 18:39:25 +0000528 return unicode("fooled you!")
Tim Peters2cfe3682001-05-05 05:36:48 +0000529 return self.it.next()
530
531 f = open(TESTFN, "w")
532 try:
533 f.write("a\n" + "b\n" + "c\n")
534 finally:
535 f.close()
536
537 f = open(TESTFN, "r")
538 # Nasty: string.join(s) can't know whether unicode.join() is needed
539 # until it's seen all of s's elements. But in this case, f's
540 # iterator cannot be restarted. So what we're testing here is
541 # whether string.join() can manage to remember everything it's seen
542 # and pass that on to unicode.join().
543 try:
544 got = " - ".join(OhPhooey(f))
Martin v. Löwis339d0f72001-08-17 18:39:25 +0000545 self.assertEqual(got, unicode("a\n - b\n - fooled you! - c\n"))
Tim Peters2cfe3682001-05-05 05:36:48 +0000546 finally:
547 f.close()
548 try:
549 unlink(TESTFN)
550 except OSError:
551 pass
Martin v. Löwis339d0f72001-08-17 18:39:25 +0000552 if not have_unicode:
553 def test_unicode_join_endcase(self): pass
Tim Peters2cfe3682001-05-05 05:36:48 +0000554
Tim Petersde9725f2001-05-05 10:06:17 +0000555 # Test iterators with 'x in y' and 'x not in y'.
556 def test_in_and_not_in(self):
Tim Peterscb8d3682001-05-05 21:05:01 +0000557 for sc5 in IteratingSequenceClass(5), SequenceClass(5):
558 for i in range(5):
559 self.assert_(i in sc5)
560 for i in "abc", -1, 5, 42.42, (3, 4), [], {1: 1}, 3-12j, sc5:
561 self.assert_(i not in sc5)
Tim Petersde9725f2001-05-05 10:06:17 +0000562
563 self.assertRaises(TypeError, lambda: 3 in 12)
564 self.assertRaises(TypeError, lambda: 3 not in map)
565
566 d = {"one": 1, "two": 2, "three": 3, 1j: 2j}
567 for k in d:
568 self.assert_(k in d)
569 self.assert_(k not in d.itervalues())
570 for v in d.values():
571 self.assert_(v in d.itervalues())
572 self.assert_(v not in d)
573 for k, v in d.iteritems():
574 self.assert_((k, v) in d.iteritems())
575 self.assert_((v, k) not in d.iteritems())
Tim Petersde9725f2001-05-05 10:06:17 +0000576
577 f = open(TESTFN, "w")
578 try:
579 f.write("a\n" "b\n" "c\n")
580 finally:
581 f.close()
582 f = open(TESTFN, "r")
583 try:
584 for chunk in "abc":
585 f.seek(0, 0)
586 self.assert_(chunk not in f)
587 f.seek(0, 0)
588 self.assert_((chunk + "\n") in f)
589 finally:
590 f.close()
591 try:
592 unlink(TESTFN)
593 except OSError:
594 pass
595
Tim Peters75f8e352001-05-05 11:33:43 +0000596 # Test iterators with operator.countOf (PySequence_Count).
597 def test_countOf(self):
598 from operator import countOf
599 self.assertEqual(countOf([1,2,2,3,2,5], 2), 3)
600 self.assertEqual(countOf((1,2,2,3,2,5), 2), 3)
601 self.assertEqual(countOf("122325", "2"), 3)
602 self.assertEqual(countOf("122325", "6"), 0)
603
604 self.assertRaises(TypeError, countOf, 42, 1)
605 self.assertRaises(TypeError, countOf, countOf, countOf)
606
607 d = {"one": 3, "two": 3, "three": 3, 1j: 2j}
608 for k in d:
609 self.assertEqual(countOf(d, k), 1)
610 self.assertEqual(countOf(d.itervalues(), 3), 3)
611 self.assertEqual(countOf(d.itervalues(), 2j), 1)
612 self.assertEqual(countOf(d.itervalues(), 1j), 0)
613
614 f = open(TESTFN, "w")
615 try:
616 f.write("a\n" "b\n" "c\n" "b\n")
617 finally:
618 f.close()
619 f = open(TESTFN, "r")
620 try:
621 for letter, count in ("a", 1), ("b", 2), ("c", 1), ("d", 0):
622 f.seek(0, 0)
623 self.assertEqual(countOf(f, letter + "\n"), count)
624 finally:
625 f.close()
626 try:
627 unlink(TESTFN)
628 except OSError:
629 pass
630
Tim Peters16a77ad2001-09-08 04:00:12 +0000631 # Test iterators with operator.indexOf (PySequence_Index).
632 def test_indexOf(self):
633 from operator import indexOf
634 self.assertEqual(indexOf([1,2,2,3,2,5], 1), 0)
635 self.assertEqual(indexOf((1,2,2,3,2,5), 2), 1)
636 self.assertEqual(indexOf((1,2,2,3,2,5), 3), 3)
637 self.assertEqual(indexOf((1,2,2,3,2,5), 5), 5)
638 self.assertRaises(ValueError, indexOf, (1,2,2,3,2,5), 0)
639 self.assertRaises(ValueError, indexOf, (1,2,2,3,2,5), 6)
640
641 self.assertEqual(indexOf("122325", "2"), 1)
642 self.assertEqual(indexOf("122325", "5"), 5)
643 self.assertRaises(ValueError, indexOf, "122325", "6")
644
645 self.assertRaises(TypeError, indexOf, 42, 1)
646 self.assertRaises(TypeError, indexOf, indexOf, indexOf)
647
648 f = open(TESTFN, "w")
649 try:
650 f.write("a\n" "b\n" "c\n" "d\n" "e\n")
651 finally:
652 f.close()
653 f = open(TESTFN, "r")
654 try:
655 fiter = iter(f)
656 self.assertEqual(indexOf(fiter, "b\n"), 1)
657 self.assertEqual(indexOf(fiter, "d\n"), 1)
658 self.assertEqual(indexOf(fiter, "e\n"), 0)
659 self.assertRaises(ValueError, indexOf, fiter, "a\n")
660 finally:
661 f.close()
662 try:
663 unlink(TESTFN)
664 except OSError:
665 pass
666
667 iclass = IteratingSequenceClass(3)
668 for i in range(3):
669 self.assertEqual(indexOf(iclass, i), i)
670 self.assertRaises(ValueError, indexOf, iclass, -1)
671
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000672 # Test iterators with file.writelines().
673 def test_writelines(self):
674 f = file(TESTFN, "w")
675
676 try:
677 self.assertRaises(TypeError, f.writelines, None)
678 self.assertRaises(TypeError, f.writelines, 42)
Tim Peters527e64f2001-10-04 05:36:56 +0000679
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000680 f.writelines(["1\n", "2\n"])
681 f.writelines(("3\n", "4\n"))
682 f.writelines({'5\n': None})
683 f.writelines({})
684
685 # Try a big chunk too.
686 class Iterator:
687 def __init__(self, start, finish):
688 self.start = start
689 self.finish = finish
690 self.i = self.start
691
692 def next(self):
693 if self.i >= self.finish:
694 raise StopIteration
695 result = str(self.i) + '\n'
696 self.i += 1
697 return result
698
699 def __iter__(self):
700 return self
701
702 class Whatever:
703 def __init__(self, start, finish):
704 self.start = start
705 self.finish = finish
706
707 def __iter__(self):
708 return Iterator(self.start, self.finish)
Tim Peters527e64f2001-10-04 05:36:56 +0000709
710 f.writelines(Whatever(6, 6+2000))
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000711 f.close()
712
713 f = file(TESTFN)
714 expected = [str(i) + "\n" for i in range(1, 2006)]
715 self.assertEqual(list(f), expected)
Tim Peters527e64f2001-10-04 05:36:56 +0000716
Tim Peters2c9aa5e2001-09-23 04:06:05 +0000717 finally:
718 f.close()
719 try:
720 unlink(TESTFN)
721 except OSError:
722 pass
723
724
Tim Petersd6d010b2001-06-21 02:49:55 +0000725 # Test iterators on RHS of unpacking assignments.
726 def test_unpack_iter(self):
727 a, b = 1, 2
728 self.assertEqual((a, b), (1, 2))
729
730 a, b, c = IteratingSequenceClass(3)
731 self.assertEqual((a, b, c), (0, 1, 2))
732
733 try: # too many values
734 a, b = IteratingSequenceClass(3)
735 except ValueError:
736 pass
737 else:
738 self.fail("should have raised ValueError")
739
740 try: # not enough values
741 a, b, c = IteratingSequenceClass(2)
742 except ValueError:
743 pass
744 else:
745 self.fail("should have raised ValueError")
746
747 try: # not iterable
748 a, b, c = len
749 except TypeError:
750 pass
751 else:
752 self.fail("should have raised TypeError")
753
754 a, b, c = {1: 42, 2: 42, 3: 42}.itervalues()
755 self.assertEqual((a, b, c), (42, 42, 42))
756
757 f = open(TESTFN, "w")
758 lines = ("a\n", "bb\n", "ccc\n")
759 try:
760 for line in lines:
761 f.write(line)
762 finally:
763 f.close()
764 f = open(TESTFN, "r")
765 try:
766 a, b, c = f
767 self.assertEqual((a, b, c), lines)
768 finally:
769 f.close()
770 try:
771 unlink(TESTFN)
772 except OSError:
773 pass
774
775 (a, b), (c,) = IteratingSequenceClass(2), {42: 24}
776 self.assertEqual((a, b, c), (0, 1, 42))
777
Guido van Rossumbb8f59a2001-12-03 19:33:25 +0000778 # Test reference count behavior
779
780 class C(object):
781 count = 0
782 def __new__(cls):
783 cls.count += 1
784 return object.__new__(cls)
785 def __del__(self):
786 cls = self.__class__
787 assert cls.count > 0
788 cls.count -= 1
789 x = C()
790 self.assertEqual(C.count, 1)
791 del x
792 self.assertEqual(C.count, 0)
793 l = [C(), C(), C()]
794 self.assertEqual(C.count, 3)
795 try:
796 a, b = iter(l)
797 except ValueError:
798 pass
799 del l
800 self.assertEqual(C.count, 0)
Fred Drake2e2be372001-09-20 21:33:42 +0000801
Guido van Rossum674eae62002-07-16 21:48:11 +0000802
803 # Make sure StopIteration is a "sink state".
804 # This tests various things that weren't sink states in Python 2.2.1,
805 # plus various things that always were fine.
806
807 def test_sinkstate_list(self):
808 # This used to fail
809 a = range(5)
810 b = iter(a)
811 self.assertEqual(list(b), range(5))
812 a.extend(range(5, 10))
813 self.assertEqual(list(b), [])
814
815 def test_sinkstate_tuple(self):
816 a = (0, 1, 2, 3, 4)
817 b = iter(a)
818 self.assertEqual(list(b), range(5))
819 self.assertEqual(list(b), [])
820
821 def test_sinkstate_string(self):
822 a = "abcde"
823 b = iter(a)
824 self.assertEqual(list(b), ['a', 'b', 'c', 'd', 'e'])
825 self.assertEqual(list(b), [])
826
827 def test_sinkstate_sequence(self):
828 # This used to fail
829 a = SequenceClass(5)
830 b = iter(a)
831 self.assertEqual(list(b), range(5))
832 a.n = 10
833 self.assertEqual(list(b), [])
834
835 def test_sinkstate_callable(self):
836 # This used to fail
837 def spam(state=[0]):
838 i = state[0]
839 state[0] = i+1
840 if i == 10:
841 raise AssertionError, "shouldn't have gotten this far"
842 return i
843 b = iter(spam, 5)
844 self.assertEqual(list(b), range(5))
845 self.assertEqual(list(b), [])
846
847 def test_sinkstate_dict(self):
848 # XXX For a more thorough test, see towards the end of:
849 # http://mail.python.org/pipermail/python-dev/2002-July/026512.html
850 a = {1:1, 2:2, 0:0, 4:4, 3:3}
851 for b in iter(a), a.iterkeys(), a.iteritems(), a.itervalues():
852 b = iter(a)
853 self.assertEqual(len(list(b)), 5)
854 self.assertEqual(list(b), [])
855
856 def test_sinkstate_yield(self):
857 def gen():
858 for i in range(5):
859 yield i
860 b = gen()
861 self.assertEqual(list(b), range(5))
862 self.assertEqual(list(b), [])
863
864 def test_sinkstate_range(self):
865 a = xrange(5)
866 b = iter(a)
867 self.assertEqual(list(b), range(5))
868 self.assertEqual(list(b), [])
869
870 def test_sinkstate_enumerate(self):
871 a = range(5)
872 e = enumerate(a)
873 b = iter(e)
874 self.assertEqual(list(b), zip(range(5), range(5)))
875 self.assertEqual(list(b), [])
876
877
Fred Drake2e2be372001-09-20 21:33:42 +0000878def test_main():
879 run_unittest(TestCase)
880
881
882if __name__ == "__main__":
883 test_main()