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.assertTrue(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 a new_style class with __iter__ but no next() method
 124     def test_new_style_iter_class(self):
 125         class IterClass(object):
 126             def __iter__(self):
 127                 return self
 128         self.assertRaises(TypeError, iter, IterClass())
 129
 130     # Test two-argument iter() with callable instance
 131     def test_iter_callable(self):
 132         class C:
 133             def __init__(self):
 134                 self.i = 0
 135             def __call__(self):
 136                 i = self.i
 137                 self.i = i + 1
 138                 if i > 100:
 139                     raise IndexError # Emergency stop
 140                 return i
 141         self.check_iterator(iter(C(), 10), range(10))
 142
 143     # Test two-argument iter() with function
 144     def test_iter_function(self):
 145         def spam(state=[0]):
 146             i = state[0]
 147             state[0] = i+1
 148             return i
 149         self.check_iterator(iter(spam, 10), range(10))
 150
 151     # Test two-argument iter() with function that raises StopIteration
 152     def test_iter_function_stop(self):
 153         def spam(state=[0]):
 154             i = state[0]
 155             if i == 10:
 156                 raise StopIteration
 157             state[0] = i+1
 158             return i
 159         self.check_iterator(iter(spam, 20), range(10))
 160
 161     # Test exception propagation through function iterator
 162     def test_exception_function(self):
 163         def spam(state=[0]):
 164             i = state[0]
 165             state[0] = i+1
 166             if i == 10:
 167                 raise RuntimeError
 168             return i
 169         res = []
 170         try:
 171             for x in iter(spam, 20):
 172                 res.append(x)
 173         except RuntimeError:
 174             self.assertEqual(res, range(10))
 175         else:
 176             self.fail("should have raised RuntimeError")
 177
 178     # Test exception propagation through sequence iterator
 179     def test_exception_sequence(self):
 180         class MySequenceClass(SequenceClass):
 181             def __getitem__(self, i):
 182                 if i == 10:
 183                     raise RuntimeError
 184                 return SequenceClass.__getitem__(self, i)
 185         res = []
 186         try:
 187             for x in MySequenceClass(20):
 188                 res.append(x)
 189         except RuntimeError:
 190             self.assertEqual(res, range(10))
 191         else:
 192             self.fail("should have raised RuntimeError")
 193
 194     # Test for StopIteration from __getitem__
 195     def test_stop_sequence(self):
 196         class MySequenceClass(SequenceClass):
 197             def __getitem__(self, i):
 198                 if i == 10:
 199                     raise StopIteration
 200                 return SequenceClass.__getitem__(self, i)
 201         self.check_for_loop(MySequenceClass(20), range(10))
 202
 203     # Test a big range
 204     def test_iter_big_range(self):
 205         self.check_for_loop(iter(range(10000)), range(10000))
 206
 207     # Test an empty list
 208     def test_iter_empty(self):
 209         self.check_for_loop(iter([]), [])
 210
 211     # Test a tuple
 212     def test_iter_tuple(self):
 213         self.check_for_loop(iter((0,1,2,3,4,5,6,7,8,9)), range(10))
 214
 215     # Test an xrange
 216     def test_iter_xrange(self):
 217         self.check_for_loop(iter(xrange(10)), range(10))
 218
 219     # Test a string
 220     def test_iter_string(self):
 221         self.check_for_loop(iter("abcde"), ["a", "b", "c", "d", "e"])
 222
 223     # Test a Unicode string
 224     if have_unicode:
 225         def test_iter_unicode(self):
 226             self.check_for_loop(iter(unicode("abcde")),
 227                                 [unicode("a"), unicode("b"), unicode("c"),
 228                                  unicode("d"), unicode("e")])
 229
 230     # Test a directory
 231     def test_iter_dict(self):
 232         dict = {}
 233         for i in range(10):
 234             dict[i] = None
 235         self.check_for_loop(dict, dict.keys())
 236
 237     # Test a file
 238     def test_iter_file(self):
 239         f = open(TESTFN, "w")
 240         try:
 241             for i in range(5):
 242                 f.write("%d\n" % i)
 243         finally:
 244             f.close()
 245         f = open(TESTFN, "r")
 246         try:
 247             self.check_for_loop(f, ["0\n", "1\n", "2\n", "3\n", "4\n"])
 248             self.check_for_loop(f, [])
 249         finally:
 250             f.close()
 251             try:
 252                 unlink(TESTFN)
 253             except OSError:
 254                 pass
 255
 256     # Test list()'s use of iterators.
 257     def test_builtin_list(self):
 258         self.assertEqual(list(SequenceClass(5)), range(5))
 259         self.assertEqual(list(SequenceClass(0)), [])
 260         self.assertEqual(list(()), [])
 261         self.assertEqual(list(range(10, -1, -1)), range(10, -1, -1))
 262
 263         d = {"one": 1, "two": 2, "three": 3}
 264         self.assertEqual(list(d), d.keys())
 265
 266         self.assertRaises(TypeError, list, list)
 267         self.assertRaises(TypeError, list, 42)
 268
 269         f = open(TESTFN, "w")
 270         try:
 271             for i in range(5):
 272                 f.write("%d\n" % i)
 273         finally:
 274             f.close()
 275         f = open(TESTFN, "r")
 276         try:
 277             self.assertEqual(list(f), ["0\n", "1\n", "2\n", "3\n", "4\n"])
 278             f.seek(0, 0)
 279             self.assertEqual(list(f),
 280                              ["0\n", "1\n", "2\n", "3\n", "4\n"])
 281         finally:
 282             f.close()
 283             try:
 284                 unlink(TESTFN)
 285             except OSError:
 286                 pass
 287
 288     # Test tuples()'s use of iterators.
 289     def test_builtin_tuple(self):
 290         self.assertEqual(tuple(SequenceClass(5)), (0, 1, 2, 3, 4))
 291         self.assertEqual(tuple(SequenceClass(0)), ())
 292         self.assertEqual(tuple([]), ())
 293         self.assertEqual(tuple(()), ())
 294         self.assertEqual(tuple("abc"), ("a", "b", "c"))
 295
 296         d = {"one": 1, "two": 2, "three": 3}
 297         self.assertEqual(tuple(d), tuple(d.keys()))
 298
 299         self.assertRaises(TypeError, tuple, list)
 300         self.assertRaises(TypeError, tuple, 42)
 301
 302         f = open(TESTFN, "w")
 303         try:
 304             for i in range(5):
 305                 f.write("%d\n" % i)
 306         finally:
 307             f.close()
 308         f = open(TESTFN, "r")
 309         try:
 310             self.assertEqual(tuple(f), ("0\n", "1\n", "2\n", "3\n", "4\n"))
 311             f.seek(0, 0)
 312             self.assertEqual(tuple(f),
 313                              ("0\n", "1\n", "2\n", "3\n", "4\n"))
 314         finally:
 315             f.close()
 316             try:
 317                 unlink(TESTFN)
 318             except OSError:
 319                 pass
 320
 321     # Test filter()'s use of iterators.
 322     def test_builtin_filter(self):
 323         self.assertEqual(filter(None, SequenceClass(5)), range(1, 5))
 324         self.assertEqual(filter(None, SequenceClass(0)), [])
 325         self.assertEqual(filter(None, ()), ())
 326         self.assertEqual(filter(None, "abc"), "abc")
 327
 328         d = {"one": 1, "two": 2, "three": 3}
 329         self.assertEqual(filter(None, d), d.keys())
 330
 331         self.assertRaises(TypeError, filter, None, list)
 332         self.assertRaises(TypeError, filter, None, 42)
 333
 334         class Boolean:
 335             def __init__(self, truth):
 336                 self.truth = truth
 337             def __nonzero__(self):
 338                 return self.truth
 339         bTrue = Boolean(1)
 340         bFalse = Boolean(0)
 341
 342         class Seq:
 343             def __init__(self, *args):
 344                 self.vals = args
 345             def __iter__(self):
 346                 class SeqIter:
 347                     def __init__(self, vals):
 348                         self.vals = vals
 349                         self.i = 0
 350                     def __iter__(self):
 351                         return self
 352                     def next(self):
 353                         i = self.i
 354                         self.i = i + 1
 355                         if i < len(self.vals):
 356                             return self.vals[i]
 357                         else:
 358                             raise StopIteration
 359                 return SeqIter(self.vals)
 360
 361         seq = Seq(*([bTrue, bFalse] * 25))
 362         self.assertEqual(filter(lambda x: not x, seq), [bFalse]*25)
 363         self.assertEqual(filter(lambda x: not x, iter(seq)), [bFalse]*25)
 364
 365     # Test max() and min()'s use of iterators.
 366     def test_builtin_max_min(self):
 367         self.assertEqual(max(SequenceClass(5)), 4)
 368         self.assertEqual(min(SequenceClass(5)), 0)
 369         self.assertEqual(max(8, -1), 8)
 370         self.assertEqual(min(8, -1), -1)
 371
 372         d = {"one": 1, "two": 2, "three": 3}
 373         self.assertEqual(max(d), "two")
 374         self.assertEqual(min(d), "one")
 375         self.assertEqual(max(d.itervalues()), 3)
 376         self.assertEqual(min(iter(d.itervalues())), 1)
 377
 378         f = open(TESTFN, "w")
 379         try:
 380             f.write("medium line\n")
 381             f.write("xtra large line\n")
 382             f.write("itty-bitty line\n")
 383         finally:
 384             f.close()
 385         f = open(TESTFN, "r")
 386         try:
 387             self.assertEqual(min(f), "itty-bitty line\n")
 388             f.seek(0, 0)
 389             self.assertEqual(max(f), "xtra large line\n")
 390         finally:
 391             f.close()
 392             try:
 393                 unlink(TESTFN)
 394             except OSError:
 395                 pass
 396
 397     # Test map()'s use of iterators.
 398     def test_builtin_map(self):
 399         self.assertEqual(map(None, SequenceClass(5)), range(5))
 400         self.assertEqual(map(lambda x: x+1, SequenceClass(5)), range(1, 6))
 401
 402         d = {"one": 1, "two": 2, "three": 3}
 403         self.assertEqual(map(None, d), d.keys())
 404         self.assertEqual(map(lambda k, d=d: (k, d[k]), d), d.items())
 405         dkeys = d.keys()
 406         expected = [(i < len(d) and dkeys[i] or None,
 407                      i,
 408                      i < len(d) and dkeys[i] or None)
 409                     for i in range(5)]
 410         self.assertEqual(map(None, d,
 411                                    SequenceClass(5),
 412                                    iter(d.iterkeys())),
 413                          expected)
 414
 415         f = open(TESTFN, "w")
 416         try:
 417             for i in range(10):
 418                 f.write("xy" * i + "\n") # line i has len 2*i+1
 419         finally:
 420             f.close()
 421         f = open(TESTFN, "r")
 422         try:
 423             self.assertEqual(map(len, f), range(1, 21, 2))
 424         finally:
 425             f.close()
 426             try:
 427                 unlink(TESTFN)
 428             except OSError:
 429                 pass
 430
 431     # Test zip()'s use of iterators.
 432     def test_builtin_zip(self):
 433         self.assertEqual(zip(), [])
 434         self.assertEqual(zip(*[]), [])
 435         self.assertEqual(zip(*[(1, 2), 'ab']), [(1, 'a'), (2, 'b')])
 436
 437         self.assertRaises(TypeError, zip, None)
 438         self.assertRaises(TypeError, zip, range(10), 42)
 439         self.assertRaises(TypeError, zip, range(10), zip)
 440
 441         self.assertEqual(zip(IteratingSequenceClass(3)),
 442                          [(0,), (1,), (2,)])
 443         self.assertEqual(zip(SequenceClass(3)),
 444                          [(0,), (1,), (2,)])
 445
 446         d = {"one": 1, "two": 2, "three": 3}
 447         self.assertEqual(d.items(), zip(d, d.itervalues()))
 448
 449         # Generate all ints starting at constructor arg.
 450         class IntsFrom:
 451             def __init__(self, start):
 452                 self.i = start
 453
 454             def __iter__(self):
 455                 return self
 456
 457             def next(self):
 458                 i = self.i
 459                 self.i = i+1
 460                 return i
 461
 462         f = open(TESTFN, "w")
 463         try:
 464             f.write("a\n" "bbb\n" "cc\n")
 465         finally:
 466             f.close()
 467         f = open(TESTFN, "r")
 468         try:
 469             self.assertEqual(zip(IntsFrom(0), f, IntsFrom(-100)),
 470                              [(0, "a\n", -100),
 471                               (1, "bbb\n", -99),
 472                               (2, "cc\n", -98)])
 473         finally:
 474             f.close()
 475             try:
 476                 unlink(TESTFN)
 477             except OSError:
 478                 pass
 479
 480         self.assertEqual(zip(xrange(5)), [(i,) for i in range(5)])
 481
 482         # Classes that lie about their lengths.
 483         class NoGuessLen5:
 484             def __getitem__(self, i):
 485                 if i >= 5:
 486                     raise IndexError
 487                 return i
 488
 489         class Guess3Len5(NoGuessLen5):
 490             def __len__(self):
 491                 return 3
 492
 493         class Guess30Len5(NoGuessLen5):
 494             def __len__(self):
 495                 return 30
 496
 497         self.assertEqual(len(Guess3Len5()), 3)
 498         self.assertEqual(len(Guess30Len5()), 30)
 499         self.assertEqual(zip(NoGuessLen5()), zip(range(5)))
 500         self.assertEqual(zip(Guess3Len5()), zip(range(5)))
 501         self.assertEqual(zip(Guess30Len5()), zip(range(5)))
 502
 503         expected = [(i, i) for i in range(5)]
 504         for x in NoGuessLen5(), Guess3Len5(), Guess30Len5():
 505             for y in NoGuessLen5(), Guess3Len5(), Guess30Len5():
 506                 self.assertEqual(zip(x, y), expected)
 507
 508     # Test reduces()'s use of iterators.
 509     def test_builtin_reduce(self):
 510         from operator import add
 511         self.assertEqual(reduce(add, SequenceClass(5)), 10)
 512         self.assertEqual(reduce(add, SequenceClass(5), 42), 52)
 513         self.assertRaises(TypeError, reduce, add, SequenceClass(0))
 514         self.assertEqual(reduce(add, SequenceClass(0), 42), 42)
 515         self.assertEqual(reduce(add, SequenceClass(1)), 0)
 516         self.assertEqual(reduce(add, SequenceClass(1), 42), 42)
 517
 518         d = {"one": 1, "two": 2, "three": 3}
 519         self.assertEqual(reduce(add, d), "".join(d.keys()))
 520
 521     # This test case will be removed if we don't have Unicode
 522     def test_unicode_join_endcase(self):
 523
 524         # This class inserts a Unicode object into its argument's natural
 525         # iteration, in the 3rd position.
 526         class OhPhooey:
 527             def __init__(self, seq):
 528                 self.it = iter(seq)
 529                 self.i = 0
 530
 531             def __iter__(self):
 532                 return self
 533
 534             def next(self):
 535                 i = self.i
 536                 self.i = i+1
 537                 if i == 2:
 538                     return unicode("fooled you!")
 539                 return self.it.next()
 540
 541         f = open(TESTFN, "w")
 542         try:
 543             f.write("a\n" + "b\n" + "c\n")
 544         finally:
 545             f.close()
 546
 547         f = open(TESTFN, "r")
 548         # Nasty:  string.join(s) can't know whether unicode.join() is needed
 549         # until it's seen all of s's elements.  But in this case, f's
 550         # iterator cannot be restarted.  So what we're testing here is
 551         # whether string.join() can manage to remember everything it's seen
 552         # and pass that on to unicode.join().
 553         try:
 554             got = " - ".join(OhPhooey(f))
 555             self.assertEqual(got, unicode("a\n - b\n - fooled you! - c\n"))
 556         finally:
 557             f.close()
 558             try:
 559                 unlink(TESTFN)
 560             except OSError:
 561                 pass
 562     if not have_unicode:
 563         def test_unicode_join_endcase(self): pass
 564
 565     # Test iterators with 'x in y' and 'x not in y'.
 566     def test_in_and_not_in(self):
 567         for sc5 in IteratingSequenceClass(5), SequenceClass(5):
 568             for i in range(5):
 569                 self.assertTrue(i in sc5)
 570             for i in "abc", -1, 5, 42.42, (3, 4), [], {1: 1}, 3-12j, sc5:
 571                 self.assertTrue(i not in sc5)
 572
 573         self.assertRaises(TypeError, lambda: 3 in 12)
 574         self.assertRaises(TypeError, lambda: 3 not in map)
 575
 576         d = {"one": 1, "two": 2, "three": 3, 1j: 2j}
 577         for k in d:
 578             self.assertTrue(k in d)
 579             self.assertTrue(k not in d.itervalues())
 580         for v in d.values():
 581             self.assertTrue(v in d.itervalues())
 582             self.assertTrue(v not in d)
 583         for k, v in d.iteritems():
 584             self.assertTrue((k, v) in d.iteritems())
 585             self.assertTrue((v, k) not in d.iteritems())
 586
 587         f = open(TESTFN, "w")
 588         try:
 589             f.write("a\n" "b\n" "c\n")
 590         finally:
 591             f.close()
 592         f = open(TESTFN, "r")
 593         try:
 594             for chunk in "abc":
 595                 f.seek(0, 0)
 596                 self.assertTrue(chunk not in f)
 597                 f.seek(0, 0)
 598                 self.assertTrue((chunk + "\n") in f)
 599         finally:
 600             f.close()
 601             try:
 602                 unlink(TESTFN)
 603             except OSError:
 604                 pass
 605
 606     # Test iterators with operator.countOf (PySequence_Count).
 607     def test_countOf(self):
 608         from operator import countOf
 609         self.assertEqual(countOf([1,2,2,3,2,5], 2), 3)
 610         self.assertEqual(countOf((1,2,2,3,2,5), 2), 3)
 611         self.assertEqual(countOf("122325", "2"), 3)
 612         self.assertEqual(countOf("122325", "6"), 0)
 613
 614         self.assertRaises(TypeError, countOf, 42, 1)
 615         self.assertRaises(TypeError, countOf, countOf, countOf)
 616
 617         d = {"one": 3, "two": 3, "three": 3, 1j: 2j}
 618         for k in d:
 619             self.assertEqual(countOf(d, k), 1)
 620         self.assertEqual(countOf(d.itervalues(), 3), 3)
 621         self.assertEqual(countOf(d.itervalues(), 2j), 1)
 622         self.assertEqual(countOf(d.itervalues(), 1j), 0)
 623
 624         f = open(TESTFN, "w")
 625         try:
 626             f.write("a\n" "b\n" "c\n" "b\n")
 627         finally:
 628             f.close()
 629         f = open(TESTFN, "r")
 630         try:
 631             for letter, count in ("a", 1), ("b", 2), ("c", 1), ("d", 0):
 632                 f.seek(0, 0)
 633                 self.assertEqual(countOf(f, letter + "\n"), count)
 634         finally:
 635             f.close()
 636             try:
 637                 unlink(TESTFN)
 638             except OSError:
 639                 pass
 640
 641     # Test iterators with operator.indexOf (PySequence_Index).
 642     def test_indexOf(self):
 643         from operator import indexOf
 644         self.assertEqual(indexOf([1,2,2,3,2,5], 1), 0)
 645         self.assertEqual(indexOf((1,2,2,3,2,5), 2), 1)
 646         self.assertEqual(indexOf((1,2,2,3,2,5), 3), 3)
 647         self.assertEqual(indexOf((1,2,2,3,2,5), 5), 5)
 648         self.assertRaises(ValueError, indexOf, (1,2,2,3,2,5), 0)
 649         self.assertRaises(ValueError, indexOf, (1,2,2,3,2,5), 6)
 650
 651         self.assertEqual(indexOf("122325", "2"), 1)
 652         self.assertEqual(indexOf("122325", "5"), 5)
 653         self.assertRaises(ValueError, indexOf, "122325", "6")
 654
 655         self.assertRaises(TypeError, indexOf, 42, 1)
 656         self.assertRaises(TypeError, indexOf, indexOf, indexOf)
 657
 658         f = open(TESTFN, "w")
 659         try:
 660             f.write("a\n" "b\n" "c\n" "d\n" "e\n")
 661         finally:
 662             f.close()
 663         f = open(TESTFN, "r")
 664         try:
 665             fiter = iter(f)
 666             self.assertEqual(indexOf(fiter, "b\n"), 1)
 667             self.assertEqual(indexOf(fiter, "d\n"), 1)
 668             self.assertEqual(indexOf(fiter, "e\n"), 0)
 669             self.assertRaises(ValueError, indexOf, fiter, "a\n")
 670         finally:
 671             f.close()
 672             try:
 673                 unlink(TESTFN)
 674             except OSError:
 675                 pass
 676
 677         iclass = IteratingSequenceClass(3)
 678         for i in range(3):
 679             self.assertEqual(indexOf(iclass, i), i)
 680         self.assertRaises(ValueError, indexOf, iclass, -1)
 681
 682     # Test iterators with file.writelines().
 683     def test_writelines(self):
 684         f = file(TESTFN, "w")
 685
 686         try:
 687             self.assertRaises(TypeError, f.writelines, None)
 688             self.assertRaises(TypeError, f.writelines, 42)
 689
 690             f.writelines(["1\n", "2\n"])
 691             f.writelines(("3\n", "4\n"))
 692             f.writelines({'5\n': None})
 693             f.writelines({})
 694
 695             # Try a big chunk too.
 696             class Iterator:
 697                 def __init__(self, start, finish):
 698                     self.start = start
 699                     self.finish = finish
 700                     self.i = self.start
 701
 702                 def next(self):
 703                     if self.i >= self.finish:
 704                         raise StopIteration
 705                     result = str(self.i) + '\n'
 706                     self.i += 1
 707                     return result
 708
 709                 def __iter__(self):
 710                     return self
 711
 712             class Whatever:
 713                 def __init__(self, start, finish):
 714                     self.start = start
 715                     self.finish = finish
 716
 717                 def __iter__(self):
 718                     return Iterator(self.start, self.finish)
 719
 720             f.writelines(Whatever(6, 6+2000))
 721             f.close()
 722
 723             f = file(TESTFN)
 724             expected = [str(i) + "\n" for i in range(1, 2006)]
 725             self.assertEqual(list(f), expected)
 726
 727         finally:
 728             f.close()
 729             try:
 730                 unlink(TESTFN)
 731             except OSError:
 732                 pass
 733
 734
 735     # Test iterators on RHS of unpacking assignments.
 736     def test_unpack_iter(self):
 737         a, b = 1, 2
 738         self.assertEqual((a, b), (1, 2))
 739
 740         a, b, c = IteratingSequenceClass(3)
 741         self.assertEqual((a, b, c), (0, 1, 2))
 742
 743         try:    # too many values
 744             a, b = IteratingSequenceClass(3)
 745         except ValueError:
 746             pass
 747         else:
 748             self.fail("should have raised ValueError")
 749
 750         try:    # not enough values
 751             a, b, c = IteratingSequenceClass(2)
 752         except ValueError:
 753             pass
 754         else:
 755             self.fail("should have raised ValueError")
 756
 757         try:    # not iterable
 758             a, b, c = len
 759         except TypeError:
 760             pass
 761         else:
 762             self.fail("should have raised TypeError")
 763
 764         a, b, c = {1: 42, 2: 42, 3: 42}.itervalues()
 765         self.assertEqual((a, b, c), (42, 42, 42))
 766
 767         f = open(TESTFN, "w")
 768         lines = ("a\n", "bb\n", "ccc\n")
 769         try:
 770             for line in lines:
 771                 f.write(line)
 772         finally:
 773             f.close()
 774         f = open(TESTFN, "r")
 775         try:
 776             a, b, c = f
 777             self.assertEqual((a, b, c), lines)
 778         finally:
 779             f.close()
 780             try:
 781                 unlink(TESTFN)
 782             except OSError:
 783                 pass
 784
 785         (a, b), (c,) = IteratingSequenceClass(2), {42: 24}
 786         self.assertEqual((a, b, c), (0, 1, 42))
 787
 788         # Test reference count behavior
 789
 790         class C(object):
 791             count = 0
 792             def __new__(cls):
 793                 cls.count += 1
 794                 return object.__new__(cls)
 795             def __del__(self):
 796                 cls = self.__class__
 797                 assert cls.count > 0
 798                 cls.count -= 1
 799         x = C()
 800         self.assertEqual(C.count, 1)
 801         del x
 802         self.assertEqual(C.count, 0)
 803         l = [C(), C(), C()]
 804         self.assertEqual(C.count, 3)
 805         try:
 806             a, b = iter(l)
 807         except ValueError:
 808             pass
 809         del l
 810         self.assertEqual(C.count, 0)
 811
 812
 813     # Make sure StopIteration is a "sink state".
 814     # This tests various things that weren't sink states in Python 2.2.1,
 815     # plus various things that always were fine.
 816
 817     def test_sinkstate_list(self):
 818         # This used to fail
 819         a = range(5)
 820         b = iter(a)
 821         self.assertEqual(list(b), range(5))
 822         a.extend(range(5, 10))
 823         self.assertEqual(list(b), [])
 824
 825     def test_sinkstate_tuple(self):
 826         a = (0, 1, 2, 3, 4)
 827         b = iter(a)
 828         self.assertEqual(list(b), range(5))
 829         self.assertEqual(list(b), [])
 830
 831     def test_sinkstate_string(self):
 832         a = "abcde"
 833         b = iter(a)
 834         self.assertEqual(list(b), ['a', 'b', 'c', 'd', 'e'])
 835         self.assertEqual(list(b), [])
 836
 837     def test_sinkstate_sequence(self):
 838         # This used to fail
 839         a = SequenceClass(5)
 840         b = iter(a)
 841         self.assertEqual(list(b), range(5))
 842         a.n = 10
 843         self.assertEqual(list(b), [])
 844
 845     def test_sinkstate_callable(self):
 846         # This used to fail
 847         def spam(state=[0]):
 848             i = state[0]
 849             state[0] = i+1
 850             if i == 10:
 851                 raise AssertionError, "shouldn't have gotten this far"
 852             return i
 853         b = iter(spam, 5)
 854         self.assertEqual(list(b), range(5))
 855         self.assertEqual(list(b), [])
 856
 857     def test_sinkstate_dict(self):
 858         # XXX For a more thorough test, see towards the end of:
 859         # http://mail.python.org/pipermail/python-dev/2002-July/026512.html
 860         a = {1:1, 2:2, 0:0, 4:4, 3:3}
 861         for b in iter(a), a.iterkeys(), a.iteritems(), a.itervalues():
 862             b = iter(a)
 863             self.assertEqual(len(list(b)), 5)
 864             self.assertEqual(list(b), [])
 865
 866     def test_sinkstate_yield(self):
 867         def gen():
 868             for i in range(5):
 869                 yield i
 870         b = gen()
 871         self.assertEqual(list(b), range(5))
 872         self.assertEqual(list(b), [])
 873
 874     def test_sinkstate_range(self):
 875         a = xrange(5)
 876         b = iter(a)
 877         self.assertEqual(list(b), range(5))
 878         self.assertEqual(list(b), [])
 879
 880     def test_sinkstate_enumerate(self):
 881         a = range(5)
 882         e = enumerate(a)
 883         b = iter(e)
 884         self.assertEqual(list(b), zip(range(5), range(5)))
 885         self.assertEqual(list(b), [])
 886
 887     def test_3720(self):
 888         # Avoid a crash, when an iterator deletes its next() method.
 889         class BadIterator(object):
 890             def __iter__(self):
 891                 return self
 892             def next(self):
 893                 del BadIterator.next
 894                 return 1
 895
 896         try:
 897             for i in BadIterator() :
 898                 pass
 899         except TypeError:
 900             pass
 901
 902
 903 def test_main():
 904     run_unittest(TestCase)
 905
 906
 907 if __name__ == "__main__":
 908     test_main()