Lib/test/test_iter.py

   1 # Test iterators.
   2
   3 import unittest
   4 from test.test_support import run_unittest, TESTFN, unlink, have_unicode
   5
   6 # Test result of triple loop (too big to inline)
   7 TRIPLETS = [(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
  21 class 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
  32 class IteratingSequenceClass:
  33     def __init__(self, n):
  34         self.n = n
  35     def __iter__(self):
  36         return BasicIterClass(self.n)
  37
  38 class 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
  49 class 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
 217     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")])
 222
 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
 249     # 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),
 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
 281     # 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),
 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
 314     # 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         bTrue = Boolean(1)
 333         bFalse = 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(*([bTrue, bFalse] * 25))
 355         self.assertEqual(filter(lambda x: not x, seq), [bFalse]*25)
 356         self.assertEqual(filter(lambda x: not x, iter(seq)), [bFalse]*25)
 357
 358     # 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
 371         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
 390     # 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())),
 406                          expected)
 407
 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))
 417         finally:
 418             f.close()
 419             try:
 420                 unlink(TESTFN)
 421             except OSError:
 422                 pass
 423
 424     # 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
 470         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
 498     # 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
 511     # This test case will be removed if we don't have Unicode
 512     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:
 528                     return unicode("fooled you!")
 529                 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))
 545             self.assertEqual(got, unicode("a\n - b\n - fooled you! - c\n"))
 546         finally:
 547             f.close()
 548             try:
 549                 unlink(TESTFN)
 550             except OSError:
 551                 pass
 552     if not have_unicode:
 553         def test_unicode_join_endcase(self): pass
 554
 555     # Test iterators with 'x in y' and 'x not in y'.
 556     def test_in_and_not_in(self):
 557         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)
 562
 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())
 576
 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
 596     # 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
 631     # 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
 672     # 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)
 679
 680             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)
 709
 710             f.writelines(Whatever(6, 6+2000))
 711             f.close()
 712
 713             f = file(TESTFN)
 714             expected = [str(i) + "\n" for i in range(1, 2006)]
 715             self.assertEqual(list(f), expected)
 716
 717         finally:
 718             f.close()
 719             try:
 720                 unlink(TESTFN)
 721             except OSError:
 722                 pass
 723
 724
 725     # 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
 778         # 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)
 801
 802
 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
 878 def test_main():
 879     run_unittest(TestCase)
 880
 881
 882 if __name__ == "__main__":
 883     test_main()