刷完这60个标准库模块，成为Python骨灰级玩家

引入

每个Pythoner对for...in这种循环都很熟悉，其使用方法如下

for i in [1,2,3,4,5]:
    print(i)

输出如下

$$\begin{gathered} 1 \\ 2 \\ 3 \\ 4 \\ 5 \end{gathered}$$

上述代码非常符合直觉，就是把$[1,2,3,4,5]$中的数据挨个挑出来，然后打印。但问题在于，for…in到底有什么能力，把$1,2,3,4,5$从列表中拿出来，for…in又怎么知道$[1,2,3,4,5]$是一个列表的？

这就涉及到了列表中的两个重要的魔术方法__iter__和__next__。

魔法函数 和for循环原理

为了说明这一点，可以做一个简单的粒子

class Test:
    def __init__(self, N):
        self.N = N
        self.index = -1
    def __iter__(self):
        return self
    def __next__(self):
        if self.index < self.N-1:
            self.index += 1
        else:
            raise StopIteration
        return self.index

上面这个类中规中矩，除了重写了两个魔法函数之外，没有任何值得关注的地方。

但是，一经for…in的调用，就会明白__iter__和__next__的含义

>>> test = Test(5)
>>> for t in test:
...   print(t)
...
0
1
2
3
4

其原理大致为

test = Test(5)
tmp = test.__iter__()
while True:
    try:
        print(test.__next__())
    except:
        break

即在进入for循环之后，先调用__iter__创建一个迭代器，然后不断地调用迭代器中的__next__函数，直到抛出StopIteration这个错误。

iter和next函数

在Python中，对__iter__和__next__这两个魔法函数，提供了外部的调用函数，即iter和next。

>>> test = Test(3)
>>> tmp = iter(test)
>>> next(tmp)
0
>>> next(tmp)
1
>>> next(tmp)
2
>>> next(tmp)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 11, in __next__
StopIteration
>>> next(test)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 11, in __next__
StopIteration

可见在弹出$0,1,2$之后，再去调用next，弹出了迭代终止的错误。

但对于列表这种数据结构，并不能直接next，只有在iter之后，才能调用

>>> x = [1,2,3]
>>> next(x)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'list' object is not an iterator
>>> y = iter(x)
>>> next(y)
1

换言之，对于一个类而言，__iter__和__next__并不需要同时具备，但一般来说__iter__返回的对象需要带有__next__魔法函数。其实现逻辑大致如下

class NextTest:
    def __init__(self, N):
        self.N = N
        self.index = -1
    def __next__(self):
        if self.index < self.N-1:
            self.index += 1
        else:
            raise StopIteration
        return self.index

class IterTest:
    def __init__(self, N):
        self.N = N
    def __iter__(self):
        return NextTest(self.N)

测试一下

test = IterTest(5)
for t in test:
    print(t)

结果为

$$\begin{gathered} 0 \\ 1 \\ 2 \\ 3 \\ 4 \end{gathered}$$

这其实也是range这个迭代器的调用逻辑。