Python 迭代器与生成器

看了书,对书中的知识进行一遍梳理加深印象。

迭代器

迭代是指重复的从对象中获取数据,直至结束。「迭代协议」概括起来就是用__iter__方法返回一个实现了__next___方法的迭代器对象。

详细:

实现了__iter__方法,表示这个目标就是可迭代类型,允许手动或者自动迭代操作。__iter__方法新建并返回一个迭代器实例,通过调用这个实例的__next__方法依次返回结果,直至抛出 StopIteration异常表示结束。

可以看到,迭代器使用了分离设计。

data(数据源) --(__iter__)--> iterator(迭代器对象) --(__next__)-->foreach

我们熟知的一些内置容器类型以及常用函数比如listrange(n)zip()都实现了迭代接口。

自定义类型

我们先来手动的实现一下迭代协议,更好的理解流程和操作方式。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class DataIter:
def __init__(self, data):
self.data = data
self.index = 0
def __next__(self):
if not self.data or self.index >= len(self.data):
raise StopIteration
d = self.data[self.index]
self.index += 1
return d
class Data:
def __init__(self, n):
self.data = list(range(n))
def __iter__(self):
return DataIter(self.data)

手工迭代:

1
2
3
4
5
6
7
8
9
10
11
12
13
>>> d = Data(2)
>>> x = d.__iter__() # x 就是迭代器实例
>>> x.__next__()
0
>>> x.__next__()
1
>>> x.__next__()
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/Users/binja/Workspaces/Python/books/python notes/ch05-iter/iter.py", line 8, in __next__
raise StopIteration
StopIteration

自动迭代:

1
2
3
>>> for i in Data(2): print(i)
0
1

辅助函数

刚才是我们自己手动实现一遍,在很多容器中已经实现了迭代接口,直接返回即可。现在使用现成的

1
2
3
4
5
6
class Data:
def __init__(self, n):
self.data = list(range(n))
def __iter__(self):
return iter(self.data)

iter就是辅助函数,还可以对序列对象以及函数、方法等课调用类型(callable)进行包装

1
2
3
4
5
6
7
8
9
10
11
12
13
>>> for i in Data(2): print(i)
...
0
1
>>> x = lambda: input("n: ")
>>> for i in iter(x, "end"): print(i)
...
n: 1
1
n: 2
2
n: end
>>>

__next__方法对应的是next函数,用于手动迭代

1
2
3
4
5
6
7
8
9
10
11
12
>>> x = iter([1, 2])
>>> x
<list_iterator object at 0x10c13fcf8>
>>> while True:
... try:
... print(next(x))
... except StopIteration:
... break
...
1
2

自动迭代

手动迭代就手动调用next函数,自动迭代就是自动化的,比如for循环,编译器会生成迭代相关指令,以实现对协议方法的调用。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
>>> def test(x):
... for i in [1, 2]: print(i)
...
>>> dis.dis(test)
2 0 SETUP_LOOP 20 (to 22)
2 LOAD_CONST 3 ((1, 2))
4 GET_ITER
>> 6 FOR_ITER 12 (to 20)
8 STORE_FAST 1 (i)
10 LOAD_GLOBAL 0 (print)
12 LOAD_FAST 1 (i)
14 CALL_FUNCTION 1
16 POP_TOP
18 JUMP_ABSOLUTE 6
>> 20 POP_BLOCK
>> 22 LOAD_CONST 0 (None)
24 RETURN_VALUE
>>>

其中4 GET_ITER就是调用__iter__返回迭代器对象(或包装)

6 FOR_ITER调用__next__返回数据

与容器对比

列表、字典等容器实现了迭代器类型。但是本质上,两者不属于同一层面。迭代器不仅是一种数据读取方法,更是一种设计模式。

  • 容器的核心是存储数据,根据数据提供操作方法。
  • 迭代器的重点是逻辑控制

生成器

生成器(generator)是迭代器的进化版本,用函数和表达式替代了接口方法。生成器函数其内部以yield返回迭代数据。这与普通函数不同,无论内部逻辑如何,函数调用总是返回生成器对象。随后以普通迭代器方式继续操作。

1
2
3
4
5
6
7
8
9
10
>>> def test():
... yield 1
... yield 2
...
>>> test()
<generator object test at 0x10c0ccfc0>
>>> for i in test(): print(i)
...
1
2

每条yield语句对应一次__next__调用。可像上面那样分为多条,也可以写在循环语句中。只要函数流程结束,就抛出迭代终止异常。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
>>> def test():
... for i in range(10):
... yield i + 100
... if i >= 2: return
...
>>> x = test()
>>> next(x)
100
>>> next(x)
101
>>> next(x)
102
>>> next(x)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
StopIteration

子迭代器

如果数据源本身就是可迭代对象,那么可使用yield from子迭代器语句。

1
2
3
4
5
6
7
8
9
10
11
12
>>> def test():
... yield from "abc"
... yield from range(3)
...
>>> for i in test(): print(i)
...
a
b
c
0
1
2

生成器表达式

与推导式唯一不同的是,生成器表达式使用小括号,而推导式使用中括号

1
2
3
>>> x = (i + 100 for i in range(7) if i % 2 == 0)
>>> x
<generator object <genexpr> at 0x10c0ccfc0>

生成器表达式可作为函数调用参数

1
2
3
4
5
6
7
8
>>> def test(x): # 当不是函数唯一参数时,不能省略小括号
... print(x)
... for i in x: print(i)
...
>>> test(i for i in range(3))
0
1
2

执行

下面来看生成器的执行过程。首先编译器为生成器函数添加标记,对于这类函数,解释器不会直接执行,而是将栈帧和代码作为参数,创建生成器实例。

1
2
3
4
5
6
>>> def test(n):
... print("gen.start")
... for i in range(n):
... print(f"gen.yield {i}")
... yield i
... print("gen.resume")

生成器对象在第一次__next__调用时触发,进入用户函数执行

1
2
3
4
5
>>> x = test(2)
>>> next(x)
gen.start
gen.yield 0
0

当执行到yield指令时,在设置好返回值后,解释器会保存线程状态,并挂起当前函数流程。至于再次调用__next__时,才恢复状态,继续执行。所以,以yield指令为切换分界线,往复交替,直到函数结束。

1
2
3
4
5
6
7
8
9
>>> next(x)
gen.resume
gen.yield 1
1
>>> next(x)
gen.resume
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
StopIteration

双向通信

生成器还有一个进化特征,就是提供双向通信能力。不仅仅是数据的提供方,还可以作为数据的接收方。生成器甚至可以在外部停止迭代,或者发送信号实现重置等自定义行为。

以上都可以通过方法send来完成,除了可以向yield发送数据以外,其余与next完全一样。需要注意的是,在发送之前,需要确保生成器已经启动。

1
2
3
4
5
6
7
8
9
10
11
>>> def test():
... while True:
... v = yield 200
... print(f"resume {v}")
...
>>> x = test()
>>> x.send(None) # 必须使用 next 或 send(None) 启动生成器
200
>>> x.send(100) # 可发送任何数据,包括 None
resume 100
200

调用close方法,解释器将终止生成器迭代

1
x.close()

该方法在生成器函数内部引发GeneratorExit异常,通知解释器结束执行。

此异常无法捕获,但是不影响finally的执行

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
>>> def test():
... for i in range(10):
... try:
... yield i
... finally:
... print("finally.")
...
>>> x = test()
>>> next(x)
0
>>> next(x)
finally.
1
>>> next(x)
finally.
2
>>> x.close()
finally.
>>> next(x)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
StopIteration

还可以像send发送数据那样,像生成器throw指定异常作为信号

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
In [10]: class ExitException(Exception): pass
In [11]: class ResetException(Exception): pass
In [12]: def test():
...: while True:
...: try:
...: v = yield
...: print(f"recv: {v}")
...: except ResetException:
...: print("reset.")
...: except ExitException:
...: print("exit.")
...: return
...:
In [13]: x = test()
In [14]: x.send(None)
In [15]: x.send(1)
recv: 1
In [16]: x.throw(ResetException)
reset.
In [17]: x.throw(ExitException)
exit.
---------------------------------------------------------------------------
StopIteration Traceback (most recent call last)
<ipython-input-17-e26b821be143> in <module>()
----> 1 x.throw(ExitException)
StopIteration:

异常属于合理流程控制,不能和错误完全等同起来。