Python学习之路34-迭代器和生成器

《流畅的Python》笔记。

本章将说明Python中迭代器和生成器的运行原理。

1. 前言

如果做严格区分,迭代器(iterator)和生成器(generator)是两个概念。迭代器是用于从集合中挨个获取元素,要求数据已存在;而生成器则是“凭空”生成元素,最典型的就是斐波那契数列。但是在Python中,大多数时候迭代器和生成器被视作同一概念。从Python2.2开始,可以使用yield关键字构建生成器,其作用和迭代器一样。在Python3中,生成器有了更广泛的用途,比如range()函数返回的就是一个类似生成器的对象,而在以前,它返回的是完整的列表。

本篇将有如下内容:

  • iter()内置函数处理可迭代对象的方式
  • 如何使用Python实现经典的迭代器模式
  • 详细说明生成器函数的工作原理
  • 如何使用生成器函数或生成器表达式代替经典的迭代器
  • 如何使用yield from语句合成生成器

2. 可迭代对象与迭代器

2.1 iter()函数

当Python解释器需要迭代对象x时,会自动调用iter(x)。内置的iter()函数的运行过程如下:

  • 检查对象是否实现了__iter__方法,如果实现了就调用它来获取一个迭代器;
  • 如果没有实现__iter__方法,但实现了__getitem__方法,Python会创建一个迭代器,尝试从索引0开始获取元素;
  • 如果上述操作都失败了,Python抛出TypeError异常,通常会提示“T object is not iterable”,其中T是目标对象所属的类。

而从上述解释可以看出,任何Python序列都可迭代的原因是,它们都实现了__getitem__方法。但iter()函数之所以要检查__getitem__方法,除了能让更多对象可迭代之外,其实还为了向下兼容。至于iter()以后还检不检查__getitem__方法就很难说了(不过目测未来很长一段时间内应该不会改变这种策略),而标准的序列类型都实现了__iter__方法,所以,如果自定义类要实现可迭代,请实现__iter__方法。

由此,我们还可得出可迭代的对象的定义:

实现了__iter__方法,能获取迭代器;或者实现了__getitem__方法,能从零开始索引的对象都是可迭代的对象。

补充

  • 从Python3.4开始,检查对象x能否迭代,最准确的方法是:调用iter(x),如果不可迭代,再处理TypeError异常。这比使用isinstance(x, abc.Iterable)更准确,因为abc.Iterable不会考虑__getitem__方法。

  • iter()函数还有一个鲜为人知的用法,即:传入两个参数,使用常规的函数或任何可调用对象创建迭代器。此时,第一个参数必须是可调用对象,第二个参数是“哨兵”。当可调用对象返回的值与“哨兵”相等时,抛弃该值,结束迭代并抛出StopIteration异常。这种用法的一个实际情况就是读取文件,当读取到空行或文件末尾时,停止读取:

    1
    2
    3
    4
    # 代码2.1
    with open("test.txt") as fp:
    for line in iter(fp.readline, "\n"):
    process_line(line)

2.2 迭代器

首先需要明确可迭代对象和迭代器之间的关系:Python从可迭代对象中获取迭代器。当对象实现了__iter__方法时,Python从它获取迭代器;当对象只实现了__getitem__方法时,Python为这个对象创建迭代器。所以,Python在迭代时始终用的是迭代器!

标准迭代器的UML继承关系图如下:

从上图以及之前的描述,我们可以总结出以下几点:

  • 具体的可迭代对象__iter__方法应该返回一个具体的迭代器
  • 具体的迭代器必须实现__next____iter__方法。__iter__方法返回迭代器本身(return self);真正的迭代操作由__next__完成,当没有可迭代元素时,它还要抛出StopIteration异常;
  • 由于迭代器也是从Iterable派生出来的,所以,迭代器是可迭代对象!

从上述内容可以猜出,应该有一个next()函数与iter()函数配对。没错,对可迭代对象的具体迭代操作就是由next()函数完成。以下是两个迭代过程:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 代码2.2
s = "ABC"
# 方法1,Python会隐式创建迭代器,并捕获StopIteration异常
for char in s:
print(char)

# 方法2,显式创建迭代器并显式迭代,此时需要手动捕获StopIteration异常
it = iter(s)
while True:
try:
print(next(it))
except StopIteration:
del it
break

如果我们要实现具体的迭代器,并不一定需要从collections.abc.Iterator继承,只需要实现__next____iter__方法即可。在Python的Lib/types.py源文件有如下注释:

1
2
3
4
# Iterators in Python aren't a matter of type but of protocol.  A large
# and changing number of builtin types implement *some* flavor of
# iterator. Don't check the type! Use hasattr to check for both
# "__iter__" and "__next__" attributes instead.

所以,这里可以给迭代器下个定义:实现了__next____iter__方法的对象就是迭代器。如果再去查看abc.Iterator的源码,可以发现如下代码:

1
2
3
4
5
6
7
8
9
10
11
12
# 代码2.3
class Iterator(Iterable):
-- snip --
@classmethod
def __subclasshook__(cls, C):
# 做了更改,实际是调用 _check_methods(C, '__iter__', '__next__')
if cls is Iterator:
if (any("__next__" in B.__dict__ for B in C.__mro__) and
any("__iter__" in B.__dict__ for B in C.__mro__)):
return True
# 希望大家看到NotImplemented能想到Python解释器后面会有什么操作
return NotImplemented # 如果猜不到,可以查看《Python学习之路32》

综上,Iterator采用的是白鹅类型技术:它实现了__subclasshook__方法,通过判断对象x是否实现了__next____iter__来判断x是否是迭代器。所以,判断对象x是否为迭代器的最好方法是调用isinstance(x, abc.Iterator)

友情提示:通过迭代器不能判断是否还有剩余的元素,迭代器也不能重置。当然,你可以为迭代器添加其他方法来实现这两种功能,但并不推荐这种做法,除非这代码只有你自己欣赏。如果想要重新迭代,请再次调用iter()函数,并传入之前的可迭代对象,传入迭代器是没有用。

2.3 典型的迭代器

下面通过实现一个Sentence类和与之配对的SentenceIterator来演示传统迭代器的实现过程:

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
# 代码2.4
import re
import reprlib

RE_WORD = re.compile("\w+")

class Sentence:
def __init__(self, text):
self.text = text
self.words = RE_WORD.findall(text)

def __iter__(self):
return SentenceIterator(self.words)

class SentenceIterator:
def __init__(self, words):
self.words = words
self.index = 0 # 保存索引

def __next__(self):
try:
word = self.words[self.index]
except IndexError: # 超出索引范围时抛出异常
raise StopIteration()
self.index += 1 # 递增索引
return word

def __iter__(self):
return self # 返回迭代器本身

这里需要指出一个典型的错误思想:把Sentence变为迭代器。迭代器是可迭代对象,但可迭代对象不能是迭代器!请不要在可迭代对象的__iter__中返回可迭代对象自身,也不要为可迭代对象添加__next__方法!这是一种常见的反模式行为。

从设计模式来讲,我们对可迭代对象并不只有逐个迭代这种方式,有可能跳跃式迭代,也有可能反向迭代。如果把一个对象设计成既是可迭代对象也是迭代器,那这个对象内部将会有成吨的if-else语句,这非常不利于维护和扩展。

3. 生成器

上述版本中的Sentence需要配备一个迭代器。而更符合Python风格的方式是用生成器函数代替SentenceIterator

3.1 生成器函数

使用生成器函数改写传统的迭代器(实际上不再定义迭代器):

1
2
3
4
5
6
# 代码3.1 Sentence中其余代码不变,且不用再定义SentenceIterator
class Sentence:
-- snip --
def __iter__(self):
for word in self.words:
yield word

解释:这里的__iter__生成器函数,调用它时会创建生成器对象,然后用这个生成器对象充当迭代器。

3.2 生成器函数工作原理

只要Python函数的定义体中有yield关键字,该函数就是生成器函数(这也是和普通函数的唯一区别)。“生成器”一词指代生成器函数,以及生成器函数构建的生成器对象,比较笼统,所以请具体语境具体分析。

生成器函数是一个生成器工厂,调用生成器函数时创建一个生成器对象,包装生成器函数的定义体

生成器对象实现了迭代器接口,通常Python会自动创建这个对象。当对生成器对象调用next()函数时,生成器函数执行到定义体中的下一个yield语句的末尾,生成yield关键字后面的表达式的值,然后停止在此处,等待下一次调用。当定义体中所有语句都执行完后,生成器函数返回,外层的生成器对象抛出StopIteration异常。

友情提醒:生成器函数并不是只执行其中的yield语句;也不是只执行到最后一个yield语句,如果最后一个yield语句后面还有代码,依然会执行。

下面是关于生成器的一个简单例子:

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
# 代码3.2
>>> def gen_AB():
... print("Start")
... yield "A"
... print("Continue")
... yield "B"
... print("End.")
...
>>> gen_AB
<function gen_AB at 0x...> # 返回值和普通函数没区别
>>> gen_AB()
<generator object gen_AB at 0x...> # 返回了一个生成器对象
>>> g = gen_AB()
>>> next(g)
Start # print("Start")
'A' # 这个是生成的值
>>> temp = next(g) # 获取生成器生成的第二个值
Continue # print("Continue")
>>> temp # 输出生成器生成的第二个值
'B' # 此时还并没有抛出异常,因为生成器函数还没执行完
>>> next(g)
End. # 生成器函数执行完毕,生成器抛出异常。
Traceback (most recent call last): # 显式调用next()需要自行捕获异常
File "<input>", line 1, in <module>
StopIteration

3.3 惰性实现与生成器表达式

上述的两个版本中,我们都用了self.words属性来保存文本中的单词,即在创建Sentence对象时就获得了所有的单词。这种方式叫做及早求值(Eager Evaluation)。而与之相反的则是惰性求值(Lazy Evaluation),通俗讲就是“等用到的时候再来求值”。及早求值可能会消耗大量内存,而惰性求值则是为了减少内存的使用。

生成器表达式以前提到过,它是用圆括号括起来的推导式(并不是生成元组)。生成器表达式可以理解为列表推导惰性版本:不会一次性构造整个列表,而是返回一个生成器,按需惰性生成元素。以下是它的一个简单示例:

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
# 代码3.3
>>> def gen_AB():
... print("Start")
... yield "A"
... print("Continue")
... yield "B"
... print("End.")
...
>>> res1 = [x * 3 for x in gen_AB()] # 这里有一个生成器,但被列表推导式全部迭代完
Start
Continue
End.
>>> res1 # 一次性生成了完整的列表
['AAA', 'BBB']
>>> res2 = (x * 3 for x in gen_AB()) # 这里其实有连个生成器
>>> res2 # 返回了一个生成器对象,并没有一次性生成所有数据,惰性
<generator object <genexpr> at 0x000001D6D34D4408>
>>> for i in res2:
... print(i)
...
Start
AAA
Continue
BBB
End.

解释:由于gen_AB()是个生成器函数,所以(x * 3 for x in gen_AB())包含了两个生成器对象,其中一个是由gen_AB()创建的,是不是有点嵌套生成器的意思?

现在我们使用re.finditer将第2版的Sentence改为惰性版本,并使用生成器表达式进一步简化代码:

1
2
3
4
5
6
7
8
9
10
# 代码3.4
class Sentence:
def __init__(self, text):
self.text = text # 去掉了self.words

def __iter__(self):
return (match.group() for match in RE_WORD.finditer(self.text))
# 不适用生成器表达式的版本如下:
# for match in RE_WORD.finditer(self.text):
# yield match.group()

友情提醒:在Python3中,如果想把某种实现变成惰性版本,一般都是可以的……

生成器表达式是创建生成器的简洁语法,这样就无需定义生成器函数,一般在情况简单时使用。不过,生成器函数灵活得多,可以使用多个语句实现更复杂的逻辑,也可以作为协程使用,还可以重用代码。

3.4 itertools模块

该模块包含了很多有用的生成器函数,这里介绍两个生成器函数itertools.countitertools.takewhile

前面介绍的生成器中的数据都是有穷集合,而itertools.count则生成无穷集合。它有两个参数起始数值start和步长stepstart默认是0step默认是1。这两个参数都支持多种数字类型,比如intfloatdecimal.Decimalfractions.Fraction。以下是它的一个示例:

1
2
3
4
5
6
7
# 代码3.5
>>> import itertools
>>> gen = itertools.count(1, 0.5)
>>> next(gen)
1
>>> next(gen)
1.5

由于itertools.count不停止生成数据,所以如果调用list(count()),你的电脑会疯狂运转,直到超出内存限制。

itertools.takewhile函数则不同,它会生成一个使用另一个生成器的生成器,在指定的函数返回False时停止。因此,这两个迭代器可以结合使用:

1
2
3
4
# 代码3.6
>>> gen = itertools.takewhile(lambda n: n < 3, itertools.count(1, 0.5))
>>> list(gen)
[1, 1.5, 2.0, 2.5]

标准库中还有很多非常有用的生成器函数,这里就不一一列出了。

3.5 yield from

如果生成器函数需要产出另一个生成器生成的值,传统的解决方法是使用嵌套for循环,比如如下函数:

1
2
3
4
5
# 代码3.7
def chain(*iterables): # iterables中的元素是可迭代对象
for it in iterables:
for i in it:
yield i

而如果使用yield from句法则可以使代码更简洁:

1
2
3
4
# 代码3.8
def chain(*iterables):
for it in iterables:
yield from it

yield from语法不仅仅是语法糖,除了代替循环之外,yield from还会创建通道,把生成器当做协程使用。

3.6 把生成器当做协程

从Python2.5起,生成器加入了一个名为.send()的方法,与.__next__方法一样,.send方法致使生成器推进到下一个yield语句。但.send方法还允许生成器的调用者向生成器传入参数,把这个参数作为对应的yield语句的返回值。这个方法让调用者和生成器之间能双向交换数据,而.__next__方法只允许调用者从生成器获取值。下面是这个方法的一个简单示例:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# 代码3.9 省略了最后抛出的StopIteration异常
>>> def test_send():
... a = yield 1
... print("At the end of function, a = ", a)
...
>>> g = test_send()
>>> next(g)
1
>>> next(g)
At the end of function, a = None # 可以看出,yield表达式是有返回值的,默认返回None
>>> g = test_send() # 新建一个生成器
>>> next(g) # 在调用send()之前,必须先至少调用过一次next()
1
>>> g.send("msg")
At the end of function, a = msg # 把我们传入的参数作为了yield表达式的返回值

这一项重要改进甚至改变了生成器的本性:像这样用的话,生成器就变为了协程。

这里是想提醒大家,请慎重使用这个方法!生成器用于生产供迭代的数据,协程是数据的消费者。为了避免不必要的麻烦,请严格区分协程和迭代,虽然协程也用到到了yield,但协程和迭代没有关系!

关于协程的内容将会在后面的文章中介绍。

4. 总结

本篇首先介绍了可迭代对象与迭代器,内容包括迭代的原理以及iter()next()函数所做的工作,然后实现了一个经典的迭代器。随后,为了让这个经典的迭代器更符合Python风格,我们讨论了生成器。这期间讲到了生成器和迭代器的关系,生成器函数及其工作原理,惰性实现和生成器表达式。根据这些内容,我们将之前传统的迭代器进行了简化。随后补充了三个内容:itertools模块中的生成器函数,yield from语法和生成器的.send()

最后,建议大家一定要多了解标准库中的生成器函数,尤其是itertools模块。

VPointer wechat
欢迎大家关注我的微信公众号"代码港"~~
您的慷慨将鼓励我继续创作~~