test_iter.py 31.5 KB
Newer Older
1 2
# Test iterators.

3
import sys
4
import unittest
5
from test.support import run_unittest, TESTFN, unlink, cpython_only
6
from test.support import check_free_after_iterating
7 8
import pickle
import collections.abc
9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28

# Test result of triple loop (too big to inline)
TRIPLETS = [(0, 0, 0), (0, 0, 1), (0, 0, 2),
            (0, 1, 0), (0, 1, 1), (0, 1, 2),
            (0, 2, 0), (0, 2, 1), (0, 2, 2),

            (1, 0, 0), (1, 0, 1), (1, 0, 2),
            (1, 1, 0), (1, 1, 1), (1, 1, 2),
            (1, 2, 0), (1, 2, 1), (1, 2, 2),

            (2, 0, 0), (2, 0, 1), (2, 0, 2),
            (2, 1, 0), (2, 1, 1), (2, 1, 2),
            (2, 2, 0), (2, 2, 1), (2, 2, 2)]

# Helper classes

class BasicIterClass:
    def __init__(self, n):
        self.n = n
        self.i = 0
29
    def __next__(self):
30 31 32 33 34
        res = self.i
        if res >= self.n:
            raise StopIteration
        self.i = res + 1
        return res
35 36
    def __iter__(self):
        return self
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52

class IteratingSequenceClass:
    def __init__(self, n):
        self.n = n
    def __iter__(self):
        return BasicIterClass(self.n)

class SequenceClass:
    def __init__(self, n):
        self.n = n
    def __getitem__(self, i):
        if 0 <= i < self.n:
            return i
        else:
            raise IndexError

53 54 55 56
class UnlimitedSequenceClass:
    def __getitem__(self, i):
        return i

57 58 59 60 61 62 63 64
class DefaultIterClass:
    pass

class NoIterClass:
    def __getitem__(self, i):
        return i
    __iter__ = None

65 66 67 68 69
# Main test suite

class TestCase(unittest.TestCase):

    # Helper to check that an iterator returns a given sequence
70 71 72
    def check_iterator(self, it, seq, pickle=True):
        if pickle:
            self.check_pickle(it, seq)
73 74 75
        res = []
        while 1:
            try:
76
                val = next(it)
77 78 79 80 81 82
            except StopIteration:
                break
            res.append(val)
        self.assertEqual(res, seq)

    # Helper to check that a for loop generates a given sequence
83 84 85
    def check_for_loop(self, expr, seq, pickle=True):
        if pickle:
            self.check_pickle(iter(expr), seq)
86 87 88 89 90
        res = []
        for val in expr:
            res.append(val)
        self.assertEqual(res, seq)

91 92
    # Helper to check picklability
    def check_pickle(self, itorg, seq):
93 94 95 96 97 98 99 100
        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
            d = pickle.dumps(itorg, proto)
            it = pickle.loads(d)
            # Cannot assert type equality because dict iterators unpickle as list
            # iterators.
            # self.assertEqual(type(itorg), type(it))
            self.assertTrue(isinstance(it, collections.abc.Iterator))
            self.assertEqual(list(it), seq)
101

102 103 104 105 106 107 108 109
            it = pickle.loads(d)
            try:
                next(it)
            except StopIteration:
                continue
            d = pickle.dumps(it, proto)
            it = pickle.loads(d)
            self.assertEqual(list(it), seq[1:])
110

111 112
    # Test basic use of iter() function
    def test_iter_basic(self):
113
        self.check_iterator(iter(range(10)), list(range(10)))
114 115 116

    # Test that iter(iter(x)) is the same as iter(x)
    def test_iter_idempotency(self):
117
        seq = list(range(10))
118 119
        it = iter(seq)
        it2 = iter(it)
120
        self.assertTrue(it is it2)
121 122 123

    # Test that for loops over iterators work
    def test_iter_for_loop(self):
124
        self.check_for_loop(iter(range(10)), list(range(10)))
125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150

    # Test several independent iterators over the same list
    def test_iter_independence(self):
        seq = range(3)
        res = []
        for i in iter(seq):
            for j in iter(seq):
                for k in iter(seq):
                    res.append((i, j, k))
        self.assertEqual(res, TRIPLETS)

    # Test triple list comprehension using iterators
    def test_nested_comprehensions_iter(self):
        seq = range(3)
        res = [(i, j, k)
               for i in iter(seq) for j in iter(seq) for k in iter(seq)]
        self.assertEqual(res, TRIPLETS)

    # Test triple list comprehension without iterators
    def test_nested_comprehensions_for(self):
        seq = range(3)
        res = [(i, j, k) for i in seq for j in seq for k in seq]
        self.assertEqual(res, TRIPLETS)

    # Test a class with __iter__ in a for loop
    def test_iter_class_for(self):
151
        self.check_for_loop(IteratingSequenceClass(10), list(range(10)))
152 153 154

    # Test a class with __iter__ with explicit iter()
    def test_iter_class_iter(self):
155
        self.check_iterator(iter(IteratingSequenceClass(10)), list(range(10)))
156 157 158

    # Test for loop on a sequence class without __iter__
    def test_seq_class_for(self):
159
        self.check_for_loop(SequenceClass(10), list(range(10)))
160 161 162

    # Test iter() on a sequence class without __iter__
    def test_seq_class_iter(self):
163
        self.check_iterator(iter(SequenceClass(10)), list(range(10)))
164

165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200
    def test_mutating_seq_class_iter_pickle(self):
        orig = SequenceClass(5)
        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
            # initial iterator
            itorig = iter(orig)
            d = pickle.dumps((itorig, orig), proto)
            it, seq = pickle.loads(d)
            seq.n = 7
            self.assertIs(type(it), type(itorig))
            self.assertEqual(list(it), list(range(7)))

            # running iterator
            next(itorig)
            d = pickle.dumps((itorig, orig), proto)
            it, seq = pickle.loads(d)
            seq.n = 7
            self.assertIs(type(it), type(itorig))
            self.assertEqual(list(it), list(range(1, 7)))

            # empty iterator
            for i in range(1, 5):
                next(itorig)
            d = pickle.dumps((itorig, orig), proto)
            it, seq = pickle.loads(d)
            seq.n = 7
            self.assertIs(type(it), type(itorig))
            self.assertEqual(list(it), list(range(5, 7)))

            # exhausted iterator
            self.assertRaises(StopIteration, next, itorig)
            d = pickle.dumps((itorig, orig), proto)
            it, seq = pickle.loads(d)
            seq.n = 7
            self.assertTrue(isinstance(it, collections.abc.Iterator))
            self.assertEqual(list(it), [])

201 202 203 204 205 206 207 208 209 210 211
    def test_mutating_seq_class_exhausted_iter(self):
        a = SequenceClass(5)
        exhit = iter(a)
        empit = iter(a)
        for x in exhit:  # exhaust the iterator
            next(empit)  # not exhausted
        a.n = 7
        self.assertEqual(list(exhit), [])
        self.assertEqual(list(empit), [5, 6])
        self.assertEqual(list(a), [0, 1, 2, 3, 4, 5, 6])

212 213 214 215 216 217 218
    # Test a new_style class with __iter__ but no next() method
    def test_new_style_iter_class(self):
        class IterClass(object):
            def __iter__(self):
                return self
        self.assertRaises(TypeError, iter, IterClass())

219 220 221 222 223 224 225 226 227 228 229
    # Test two-argument iter() with callable instance
    def test_iter_callable(self):
        class C:
            def __init__(self):
                self.i = 0
            def __call__(self):
                i = self.i
                self.i = i + 1
                if i > 100:
                    raise IndexError # Emergency stop
                return i
230
        self.check_iterator(iter(C(), 10), list(range(10)), pickle=False)
231 232 233 234 235 236 237

    # Test two-argument iter() with function
    def test_iter_function(self):
        def spam(state=[0]):
            i = state[0]
            state[0] = i+1
            return i
238
        self.check_iterator(iter(spam, 10), list(range(10)), pickle=False)
239 240 241 242 243 244 245 246 247

    # Test two-argument iter() with function that raises StopIteration
    def test_iter_function_stop(self):
        def spam(state=[0]):
            i = state[0]
            if i == 10:
                raise StopIteration
            state[0] = i+1
            return i
248
        self.check_iterator(iter(spam, 20), list(range(10)), pickle=False)
249 250 251 252 253 254 255 256 257 258 259 260 261 262

    # Test exception propagation through function iterator
    def test_exception_function(self):
        def spam(state=[0]):
            i = state[0]
            state[0] = i+1
            if i == 10:
                raise RuntimeError
            return i
        res = []
        try:
            for x in iter(spam, 20):
                res.append(x)
        except RuntimeError:
263
            self.assertEqual(res, list(range(10)))
264 265 266 267 268 269 270 271 272 273 274 275 276 277 278
        else:
            self.fail("should have raised RuntimeError")

    # Test exception propagation through sequence iterator
    def test_exception_sequence(self):
        class MySequenceClass(SequenceClass):
            def __getitem__(self, i):
                if i == 10:
                    raise RuntimeError
                return SequenceClass.__getitem__(self, i)
        res = []
        try:
            for x in MySequenceClass(20):
                res.append(x)
        except RuntimeError:
279
            self.assertEqual(res, list(range(10)))
280 281 282 283 284 285 286 287 288 289
        else:
            self.fail("should have raised RuntimeError")

    # Test for StopIteration from __getitem__
    def test_stop_sequence(self):
        class MySequenceClass(SequenceClass):
            def __getitem__(self, i):
                if i == 10:
                    raise StopIteration
                return SequenceClass.__getitem__(self, i)
290
        self.check_for_loop(MySequenceClass(20), list(range(10)), pickle=False)
291 292 293

    # Test a big range
    def test_iter_big_range(self):
294
        self.check_for_loop(iter(range(10000)), list(range(10000)))
295 296 297 298 299 300 301

    # Test an empty list
    def test_iter_empty(self):
        self.check_for_loop(iter([]), [])

    # Test a tuple
    def test_iter_tuple(self):
302
        self.check_for_loop(iter((0,1,2,3,4,5,6,7,8,9)), list(range(10)))
303

304 305 306
    # Test a range
    def test_iter_range(self):
        self.check_for_loop(iter(range(10)), list(range(10)))
307 308 309 310 311 312 313 314 315 316

    # Test a string
    def test_iter_string(self):
        self.check_for_loop(iter("abcde"), ["a", "b", "c", "d", "e"])

    # Test a directory
    def test_iter_dict(self):
        dict = {}
        for i in range(10):
            dict[i] = None
317
        self.check_for_loop(dict, list(dict.keys()))
318 319 320 321 322 323 324 325 326 327 328

    # Test a file
    def test_iter_file(self):
        f = open(TESTFN, "w")
        try:
            for i in range(5):
                f.write("%d\n" % i)
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
329 330
            self.check_for_loop(f, ["0\n", "1\n", "2\n", "3\n", "4\n"], pickle=False)
            self.check_for_loop(f, [], pickle=False)
331 332 333 334 335 336 337
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

338 339
    # Test list()'s use of iterators.
    def test_builtin_list(self):
340
        self.assertEqual(list(SequenceClass(5)), list(range(5)))
341 342 343 344
        self.assertEqual(list(SequenceClass(0)), [])
        self.assertEqual(list(()), [])

        d = {"one": 1, "two": 2, "three": 3}
345
        self.assertEqual(list(d), list(d.keys()))
346 347 348 349 350 351 352 353 354 355 356 357 358 359

        self.assertRaises(TypeError, list, list)
        self.assertRaises(TypeError, list, 42)

        f = open(TESTFN, "w")
        try:
            for i in range(5):
                f.write("%d\n" % i)
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
            self.assertEqual(list(f), ["0\n", "1\n", "2\n", "3\n", "4\n"])
            f.seek(0, 0)
360
            self.assertEqual(list(f),
361 362 363 364 365 366 367 368
                             ["0\n", "1\n", "2\n", "3\n", "4\n"])
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392
    # Test tuples()'s use of iterators.
    def test_builtin_tuple(self):
        self.assertEqual(tuple(SequenceClass(5)), (0, 1, 2, 3, 4))
        self.assertEqual(tuple(SequenceClass(0)), ())
        self.assertEqual(tuple([]), ())
        self.assertEqual(tuple(()), ())
        self.assertEqual(tuple("abc"), ("a", "b", "c"))

        d = {"one": 1, "two": 2, "three": 3}
        self.assertEqual(tuple(d), tuple(d.keys()))

        self.assertRaises(TypeError, tuple, list)
        self.assertRaises(TypeError, tuple, 42)

        f = open(TESTFN, "w")
        try:
            for i in range(5):
                f.write("%d\n" % i)
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
            self.assertEqual(tuple(f), ("0\n", "1\n", "2\n", "3\n", "4\n"))
            f.seek(0, 0)
393
            self.assertEqual(tuple(f),
394 395 396 397 398 399 400 401
                             ("0\n", "1\n", "2\n", "3\n", "4\n"))
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

402 403
    # Test filter()'s use of iterators.
    def test_builtin_filter(self):
404 405 406 407 408
        self.assertEqual(list(filter(None, SequenceClass(5))),
                         list(range(1, 5)))
        self.assertEqual(list(filter(None, SequenceClass(0))), [])
        self.assertEqual(list(filter(None, ())), [])
        self.assertEqual(list(filter(None, "abc")), ["a", "b", "c"])
409 410

        d = {"one": 1, "two": 2, "three": 3}
411
        self.assertEqual(list(filter(None, d)), list(d.keys()))
412 413 414 415 416 417 418

        self.assertRaises(TypeError, filter, None, list)
        self.assertRaises(TypeError, filter, None, 42)

        class Boolean:
            def __init__(self, truth):
                self.truth = truth
419
            def __bool__(self):
420
                return self.truth
421 422
        bTrue = Boolean(True)
        bFalse = Boolean(False)
423 424 425 426 427 428 429 430 431 432 433

        class Seq:
            def __init__(self, *args):
                self.vals = args
            def __iter__(self):
                class SeqIter:
                    def __init__(self, vals):
                        self.vals = vals
                        self.i = 0
                    def __iter__(self):
                        return self
434
                    def __next__(self):
435 436 437 438 439 440 441 442
                        i = self.i
                        self.i = i + 1
                        if i < len(self.vals):
                            return self.vals[i]
                        else:
                            raise StopIteration
                return SeqIter(self.vals)

443
        seq = Seq(*([bTrue, bFalse] * 25))
444 445
        self.assertEqual(list(filter(lambda x: not x, seq)), [bFalse]*25)
        self.assertEqual(list(filter(lambda x: not x, iter(seq))), [bFalse]*25)
446

447 448 449 450 451 452 453 454 455 456
    # Test max() and min()'s use of iterators.
    def test_builtin_max_min(self):
        self.assertEqual(max(SequenceClass(5)), 4)
        self.assertEqual(min(SequenceClass(5)), 0)
        self.assertEqual(max(8, -1), 8)
        self.assertEqual(min(8, -1), -1)

        d = {"one": 1, "two": 2, "three": 3}
        self.assertEqual(max(d), "two")
        self.assertEqual(min(d), "one")
457 458
        self.assertEqual(max(d.values()), 3)
        self.assertEqual(min(iter(d.values())), 1)
459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478

        f = open(TESTFN, "w")
        try:
            f.write("medium line\n")
            f.write("xtra large line\n")
            f.write("itty-bitty line\n")
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
            self.assertEqual(min(f), "itty-bitty line\n")
            f.seek(0, 0)
            self.assertEqual(max(f), "xtra large line\n")
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

479 480
    # Test map()'s use of iterators.
    def test_builtin_map(self):
481 482
        self.assertEqual(list(map(lambda x: x+1, SequenceClass(5))),
                         list(range(1, 6)))
483 484

        d = {"one": 1, "two": 2, "three": 3}
485 486
        self.assertEqual(list(map(lambda k, d=d: (k, d[k]), d)),
                         list(d.items()))
487
        dkeys = list(d.keys())
488 489 490
        expected = [(i < len(d) and dkeys[i] or None,
                     i,
                     i < len(d) and dkeys[i] or None)
491
                    for i in range(3)]
492 493 494 495 496 497 498 499 500

        f = open(TESTFN, "w")
        try:
            for i in range(10):
                f.write("xy" * i + "\n") # line i has len 2*i+1
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
501
            self.assertEqual(list(map(len, f)), list(range(1, 21, 2)))
502 503 504 505 506 507 508
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

509 510
    # Test zip()'s use of iterators.
    def test_builtin_zip(self):
511 512 513
        self.assertEqual(list(zip()), [])
        self.assertEqual(list(zip(*[])), [])
        self.assertEqual(list(zip(*[(1, 2), 'ab'])), [(1, 'a'), (2, 'b')])
514

515 516 517 518
        self.assertRaises(TypeError, zip, None)
        self.assertRaises(TypeError, zip, range(10), 42)
        self.assertRaises(TypeError, zip, range(10), zip)

519
        self.assertEqual(list(zip(IteratingSequenceClass(3))),
520
                         [(0,), (1,), (2,)])
521
        self.assertEqual(list(zip(SequenceClass(3))),
522 523 524
                         [(0,), (1,), (2,)])

        d = {"one": 1, "two": 2, "three": 3}
525
        self.assertEqual(list(d.items()), list(zip(d, d.values())))
526 527 528 529 530 531 532 533 534

        # Generate all ints starting at constructor arg.
        class IntsFrom:
            def __init__(self, start):
                self.i = start

            def __iter__(self):
                return self

535
            def __next__(self):
536 537 538 539 540 541 542 543 544 545 546
                i = self.i
                self.i = i+1
                return i

        f = open(TESTFN, "w")
        try:
            f.write("a\n" "bbb\n" "cc\n")
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
547
            self.assertEqual(list(zip(IntsFrom(0), f, IntsFrom(-100))),
548 549 550 551 552 553 554 555 556 557
                             [(0, "a\n", -100),
                              (1, "bbb\n", -99),
                              (2, "cc\n", -98)])
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

558
        self.assertEqual(list(zip(range(5))), [(i,) for i in range(5)])
559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574

        # Classes that lie about their lengths.
        class NoGuessLen5:
            def __getitem__(self, i):
                if i >= 5:
                    raise IndexError
                return i

        class Guess3Len5(NoGuessLen5):
            def __len__(self):
                return 3

        class Guess30Len5(NoGuessLen5):
            def __len__(self):
                return 30

575 576 577
        def lzip(*args):
            return list(zip(*args))

578 579
        self.assertEqual(len(Guess3Len5()), 3)
        self.assertEqual(len(Guess30Len5()), 30)
580 581 582
        self.assertEqual(lzip(NoGuessLen5()), lzip(range(5)))
        self.assertEqual(lzip(Guess3Len5()), lzip(range(5)))
        self.assertEqual(lzip(Guess30Len5()), lzip(range(5)))
583 584 585 586

        expected = [(i, i) for i in range(5)]
        for x in NoGuessLen5(), Guess3Len5(), Guess30Len5():
            for y in NoGuessLen5(), Guess3Len5(), Guess30Len5():
587
                self.assertEqual(lzip(x, y), expected)
588

589 590 591 592 593 594 595 596 597 598 599 600
    def test_unicode_join_endcase(self):

        # This class inserts a Unicode object into its argument's natural
        # iteration, in the 3rd position.
        class OhPhooey:
            def __init__(self, seq):
                self.it = iter(seq)
                self.i = 0

            def __iter__(self):
                return self

601
            def __next__(self):
602 603 604
                i = self.i
                self.i = i+1
                if i == 2:
605
                    return "fooled you!"
606
                return next(self.it)
607 608 609 610 611 612 613 614 615 616 617 618 619 620 621

        f = open(TESTFN, "w")
        try:
            f.write("a\n" + "b\n" + "c\n")
        finally:
            f.close()

        f = open(TESTFN, "r")
        # Nasty:  string.join(s) can't know whether unicode.join() is needed
        # until it's seen all of s's elements.  But in this case, f's
        # iterator cannot be restarted.  So what we're testing here is
        # whether string.join() can manage to remember everything it's seen
        # and pass that on to unicode.join().
        try:
            got = " - ".join(OhPhooey(f))
622
            self.assertEqual(got, "a\n - b\n - fooled you! - c\n")
623 624 625 626 627 628 629
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

630 631
    # Test iterators with 'x in y' and 'x not in y'.
    def test_in_and_not_in(self):
632 633
        for sc5 in IteratingSequenceClass(5), SequenceClass(5):
            for i in range(5):
634
                self.assertIn(i, sc5)
635
            for i in "abc", -1, 5, 42.42, (3, 4), [], {1: 1}, 3-12j, sc5:
636
                self.assertNotIn(i, sc5)
637 638 639 640 641 642

        self.assertRaises(TypeError, lambda: 3 in 12)
        self.assertRaises(TypeError, lambda: 3 not in map)

        d = {"one": 1, "two": 2, "three": 3, 1j: 2j}
        for k in d:
643 644
            self.assertIn(k, d)
            self.assertNotIn(k, d.values())
645
        for v in d.values():
646 647
            self.assertIn(v, d.values())
            self.assertNotIn(v, d)
648
        for k, v in d.items():
649 650
            self.assertIn((k, v), d.items())
            self.assertNotIn((v, k), d.items())
651 652 653 654 655 656 657 658 659 660

        f = open(TESTFN, "w")
        try:
            f.write("a\n" "b\n" "c\n")
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
            for chunk in "abc":
                f.seek(0, 0)
661
                self.assertNotIn(chunk, f)
662
                f.seek(0, 0)
663
                self.assertIn((chunk + "\n"), f)
664 665 666 667 668 669 670
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

671 672 673 674 675 676 677 678 679 680 681 682 683 684
    # Test iterators with operator.countOf (PySequence_Count).
    def test_countOf(self):
        from operator import countOf
        self.assertEqual(countOf([1,2,2,3,2,5], 2), 3)
        self.assertEqual(countOf((1,2,2,3,2,5), 2), 3)
        self.assertEqual(countOf("122325", "2"), 3)
        self.assertEqual(countOf("122325", "6"), 0)

        self.assertRaises(TypeError, countOf, 42, 1)
        self.assertRaises(TypeError, countOf, countOf, countOf)

        d = {"one": 3, "two": 3, "three": 3, 1j: 2j}
        for k in d:
            self.assertEqual(countOf(d, k), 1)
685 686 687
        self.assertEqual(countOf(d.values(), 3), 3)
        self.assertEqual(countOf(d.values(), 2j), 1)
        self.assertEqual(countOf(d.values(), 1j), 0)
688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705

        f = open(TESTFN, "w")
        try:
            f.write("a\n" "b\n" "c\n" "b\n")
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
            for letter, count in ("a", 1), ("b", 2), ("c", 1), ("d", 0):
                f.seek(0, 0)
                self.assertEqual(countOf(f, letter + "\n"), count)
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746
    # Test iterators with operator.indexOf (PySequence_Index).
    def test_indexOf(self):
        from operator import indexOf
        self.assertEqual(indexOf([1,2,2,3,2,5], 1), 0)
        self.assertEqual(indexOf((1,2,2,3,2,5), 2), 1)
        self.assertEqual(indexOf((1,2,2,3,2,5), 3), 3)
        self.assertEqual(indexOf((1,2,2,3,2,5), 5), 5)
        self.assertRaises(ValueError, indexOf, (1,2,2,3,2,5), 0)
        self.assertRaises(ValueError, indexOf, (1,2,2,3,2,5), 6)

        self.assertEqual(indexOf("122325", "2"), 1)
        self.assertEqual(indexOf("122325", "5"), 5)
        self.assertRaises(ValueError, indexOf, "122325", "6")

        self.assertRaises(TypeError, indexOf, 42, 1)
        self.assertRaises(TypeError, indexOf, indexOf, indexOf)

        f = open(TESTFN, "w")
        try:
            f.write("a\n" "b\n" "c\n" "d\n" "e\n")
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
            fiter = iter(f)
            self.assertEqual(indexOf(fiter, "b\n"), 1)
            self.assertEqual(indexOf(fiter, "d\n"), 1)
            self.assertEqual(indexOf(fiter, "e\n"), 0)
            self.assertRaises(ValueError, indexOf, fiter, "a\n")
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

        iclass = IteratingSequenceClass(3)
        for i in range(3):
            self.assertEqual(indexOf(iclass, i), i)
        self.assertRaises(ValueError, indexOf, iclass, -1)

747 748
    # Test iterators with file.writelines().
    def test_writelines(self):
749
        f = open(TESTFN, "w")
750 751 752 753

        try:
            self.assertRaises(TypeError, f.writelines, None)
            self.assertRaises(TypeError, f.writelines, 42)
Tim Peters's avatar
Tim Peters committed
754

755 756 757 758 759 760 761 762 763 764 765 766
            f.writelines(["1\n", "2\n"])
            f.writelines(("3\n", "4\n"))
            f.writelines({'5\n': None})
            f.writelines({})

            # Try a big chunk too.
            class Iterator:
                def __init__(self, start, finish):
                    self.start = start
                    self.finish = finish
                    self.i = self.start

767
                def __next__(self):
768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783
                    if self.i >= self.finish:
                        raise StopIteration
                    result = str(self.i) + '\n'
                    self.i += 1
                    return result

                def __iter__(self):
                    return self

            class Whatever:
                def __init__(self, start, finish):
                    self.start = start
                    self.finish = finish

                def __iter__(self):
                    return Iterator(self.start, self.finish)
Tim Peters's avatar
Tim Peters committed
784 785

            f.writelines(Whatever(6, 6+2000))
786 787
            f.close()

788
            f = open(TESTFN)
789 790
            expected = [str(i) + "\n" for i in range(1, 2006)]
            self.assertEqual(list(f), expected)
Tim Peters's avatar
Tim Peters committed
791

792 793 794 795 796 797 798 799
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass


800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828
    # Test iterators on RHS of unpacking assignments.
    def test_unpack_iter(self):
        a, b = 1, 2
        self.assertEqual((a, b), (1, 2))

        a, b, c = IteratingSequenceClass(3)
        self.assertEqual((a, b, c), (0, 1, 2))

        try:    # too many values
            a, b = IteratingSequenceClass(3)
        except ValueError:
            pass
        else:
            self.fail("should have raised ValueError")

        try:    # not enough values
            a, b, c = IteratingSequenceClass(2)
        except ValueError:
            pass
        else:
            self.fail("should have raised ValueError")

        try:    # not iterable
            a, b, c = len
        except TypeError:
            pass
        else:
            self.fail("should have raised TypeError")

829
        a, b, c = {1: 42, 2: 42, 3: 42}.values()
830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852
        self.assertEqual((a, b, c), (42, 42, 42))

        f = open(TESTFN, "w")
        lines = ("a\n", "bb\n", "ccc\n")
        try:
            for line in lines:
                f.write(line)
        finally:
            f.close()
        f = open(TESTFN, "r")
        try:
            a, b, c = f
            self.assertEqual((a, b, c), lines)
        finally:
            f.close()
            try:
                unlink(TESTFN)
            except OSError:
                pass

        (a, b), (c,) = IteratingSequenceClass(2), {42: 24}
        self.assertEqual((a, b, c), (0, 1, 42))

853

854 855
    @cpython_only
    def test_ref_counting_behavior(self):
856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876
        class C(object):
            count = 0
            def __new__(cls):
                cls.count += 1
                return object.__new__(cls)
            def __del__(self):
                cls = self.__class__
                assert cls.count > 0
                cls.count -= 1
        x = C()
        self.assertEqual(C.count, 1)
        del x
        self.assertEqual(C.count, 0)
        l = [C(), C(), C()]
        self.assertEqual(C.count, 3)
        try:
            a, b = iter(l)
        except ValueError:
            pass
        del l
        self.assertEqual(C.count, 0)
877

878 879 880 881 882 883 884

    # Make sure StopIteration is a "sink state".
    # This tests various things that weren't sink states in Python 2.2.1,
    # plus various things that always were fine.

    def test_sinkstate_list(self):
        # This used to fail
885
        a = list(range(5))
886
        b = iter(a)
887
        self.assertEqual(list(b), list(range(5)))
888 889 890 891 892 893
        a.extend(range(5, 10))
        self.assertEqual(list(b), [])

    def test_sinkstate_tuple(self):
        a = (0, 1, 2, 3, 4)
        b = iter(a)
894
        self.assertEqual(list(b), list(range(5)))
895 896 897 898 899 900 901 902 903 904 905 906
        self.assertEqual(list(b), [])

    def test_sinkstate_string(self):
        a = "abcde"
        b = iter(a)
        self.assertEqual(list(b), ['a', 'b', 'c', 'd', 'e'])
        self.assertEqual(list(b), [])

    def test_sinkstate_sequence(self):
        # This used to fail
        a = SequenceClass(5)
        b = iter(a)
907
        self.assertEqual(list(b), list(range(5)))
908 909 910 911 912 913 914 915 916
        a.n = 10
        self.assertEqual(list(b), [])

    def test_sinkstate_callable(self):
        # This used to fail
        def spam(state=[0]):
            i = state[0]
            state[0] = i+1
            if i == 10:
917
                raise AssertionError("shouldn't have gotten this far")
918 919
            return i
        b = iter(spam, 5)
920
        self.assertEqual(list(b), list(range(5)))
921 922 923 924 925 926
        self.assertEqual(list(b), [])

    def test_sinkstate_dict(self):
        # XXX For a more thorough test, see towards the end of:
        # http://mail.python.org/pipermail/python-dev/2002-July/026512.html
        a = {1:1, 2:2, 0:0, 4:4, 3:3}
927
        for b in iter(a), a.keys(), a.items(), a.values():
928 929 930 931 932 933 934 935 936
            b = iter(a)
            self.assertEqual(len(list(b)), 5)
            self.assertEqual(list(b), [])

    def test_sinkstate_yield(self):
        def gen():
            for i in range(5):
                yield i
        b = gen()
937
        self.assertEqual(list(b), list(range(5)))
938 939 940
        self.assertEqual(list(b), [])

    def test_sinkstate_range(self):
941
        a = range(5)
942
        b = iter(a)
943
        self.assertEqual(list(b), list(range(5)))
944 945 946 947 948 949
        self.assertEqual(list(b), [])

    def test_sinkstate_enumerate(self):
        a = range(5)
        e = enumerate(a)
        b = iter(e)
950
        self.assertEqual(list(b), list(zip(range(5), range(5))))
951 952
        self.assertEqual(list(b), [])

953 954 955 956 957 958 959 960 961 962 963 964 965 966 967
    def test_3720(self):
        # Avoid a crash, when an iterator deletes its next() method.
        class BadIterator(object):
            def __iter__(self):
                return self
            def __next__(self):
                del BadIterator.__next__
                return 1

        try:
            for i in BadIterator() :
                pass
        except TypeError:
            pass

Ezio Melotti's avatar
Ezio Melotti committed
968 969 970 971 972 973 974 975 976 977 978 979 980 981 982
    def test_extending_list_with_iterator_does_not_segfault(self):
        # The code to extend a list with an iterator has a fair
        # amount of nontrivial logic in terms of guessing how
        # much memory to allocate in advance, "stealing" refs,
        # and then shrinking at the end.  This is a basic smoke
        # test for that scenario.
        def gen():
            for i in range(500):
                yield i
        lst = [0] * 500
        for i in range(240):
            lst.pop(0)
        lst.extend(gen())
        self.assertEqual(len(lst), 760)

983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002
    @cpython_only
    def test_iter_overflow(self):
        # Test for the issue 22939
        it = iter(UnlimitedSequenceClass())
        # Manually set `it_index` to PY_SSIZE_T_MAX-2 without a loop
        it.__setstate__(sys.maxsize - 2)
        self.assertEqual(next(it), sys.maxsize - 2)
        self.assertEqual(next(it), sys.maxsize - 1)
        with self.assertRaises(OverflowError):
            next(it)
        # Check that Overflow error is always raised
        with self.assertRaises(OverflowError):
            next(it)

    def test_iter_neg_setstate(self):
        it = iter(UnlimitedSequenceClass())
        it.__setstate__(-42)
        self.assertEqual(next(it), 0)
        self.assertEqual(next(it), 1)

1003 1004 1005
    def test_free_after_iterating(self):
        check_free_after_iterating(self, iter, SequenceClass, (0,))

1006 1007 1008 1009
    def test_error_iter(self):
        for typ in (DefaultIterClass, NoIterClass):
            self.assertRaises(TypeError, iter, typ())

1010

1011 1012 1013 1014 1015 1016
def test_main():
    run_unittest(TestCase)


if __name__ == "__main__":
    test_main()