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异常。这种用法的一个实际情况就是读取文件，当读取到空行或文件末尾时，停止读取：

  # 代码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()函数完成。以下是两个迭代过程：

# 代码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源文件有如下注释：

# 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的源码，可以发现如下代码：

# 代码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来演示传统迭代器的实现过程：

# 代码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 生成器函数

使用生成器函数改写传统的迭代器（实际上不再定义迭代器）：

# 代码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语句后面还有代码，依然会执行。

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

# 代码3.2
>>> def gen_AB():
...     print("Start")
...     yield "A"
...     print("Continue")
...     yield "B"
...     print("End.")
...
>>> gen_AB
0x...>   # 返回值和普通函数没区别
>>> gen_AB()
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 "", line 1, in 
StopIteration
复制代码

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

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

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

# 代码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   # 返回了一个生成器对象，并没有一次性生成所有数据，惰性
 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改为惰性版本，并使用生成器表达式进一步简化代码：

# 代码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.count和itertools.takewhile。

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

# 代码3.5
>>> import itertools
>>> gen = itertools.count(1, 0.5)
>>> next(gen)
1
>>> next(gen)
1.5
复制代码

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

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

# 代码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循环，比如如下函数：

# 代码3.7
def chain(*iterables):  # iterables中的元素是可迭代对象
    for it in iterables:
        for i in it:
            yield i
复制代码

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

# 代码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__方法只允许调用者从生成器获取值。下面是这个方法的一个简单示例：

# 代码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模块。

迎大家关注我的微信公众号"代码港" & 个人网站 www.vpointer.net ~