迭代器协议

for 背后到底发生了什么

for x in [1, 2, 3]:
    print(x)

Python 实际上做的事：

it = iter([1, 2, 3])         # 拿到迭代器
while True:
    try:
        x = next(it)           # 取下一个
    except StopIteration:
        break
    print(x)

迭代器协议就是两个魔法方法：__iter__ 和 __next__。

自己写一个迭代器

class Countdown:
    def __init__(self, start):
        self.n = start

    def __iter__(self):
        return self            # 我自己就是迭代器

    def __next__(self):
        if self.n <= 0:
            raise StopIteration
        v = self.n
        self.n -= 1
        return v

for x in Countdown(3):
    print(x)
# 3 2 1

迭代器是"一次性"的

c = Countdown(3)
list(c)        # [3, 2, 1]
list(c)        # [] —— 已经用完了

要再用一次就得重新创建。这和列表（可迭代对象）不同——列表可以反复迭代。

可迭代对象 vs 迭代器

lst = [1, 2, 3]
iter(lst) is lst    # False —— list 本身不是迭代器
it = iter(lst)
iter(it) is it      # True —— iter 自己就是迭代器

自定义可迭代对象（更好）

让 __iter__ 每次返回新迭代器，就能反复迭代：

class CountdownIterable:
    def __init__(self, start):
        self.start = start

    def __iter__(self):
        return Countdown(self.start)    # 每次新建

c = CountdownIterable(3)
list(c)        # [3, 2, 1]
list(c)        # [3, 2, 1]   还能再来

实战：链式迭代器

class Cycle:
    def __init__(self, items):
        self.items = items
        self.i = 0

    def __iter__(self):
        return self

    def __next__(self):
        v = self.items[self.i % len(self.items)]
        self.i += 1
        return v

c = Cycle(["A", "B", "C"])
for _ in range(7):
    print(next(c), end=" ")
# A B C A B C A

itertools 全家桶

标准库已经写好了一堆迭代器工具：

from itertools import count, cycle, chain, islice

list(islice(count(10, 2), 5))    # [10, 12, 14, 16, 18]
list(islice(cycle("AB"), 5))     # ['A', 'B', 'A', 'B', 'A']
list(chain([1,2], [3,4]))        # [1, 2, 3, 4]

下一篇讲生成器与 yield——写迭代器的更优雅方式。