Python 迭代器技巧|极客教程

Python 迭代器技巧，与许多其他编程语言相比，我喜欢美丽而清晰的Python语法。例如低调的for-in循环，Python的美从中得以展现出来，读起来就像英文句子那么自然：

numbers = [1, 2, 3]
for n in numbers:
    print(n)

但这种优雅的循环结构在Python内部是如何工作的？循环如何从正在循环的对象中获取单个元素？如何在自己的Python对象中支持这种编程风格？

答案是在Python中使用迭代器协议，只要对象支持__iter__和__next__双下划线方法，那么就能使用for-in循环。

与装饰器一样，迭代器及相关技术乍一看可能显得非常神秘和复杂，所以这里分阶段逐步介绍。

本节中将编写几个支持迭代器协议的Python类，这些示例和测试实现浅显易懂，你可以参照它们来加深对迭代器的理解。

首先关注Python 3中迭代器的核心机制，但这里会避免牵扯其他无关内容，以便清楚地介绍迭代器的基本行为。

每个例子最后都会用for-in循环再次实现。本节最后还将讨论迭代器在Python 2和Python 3之间的差异。
准备好了吗？让我们开始吧！

Python 迭代器技巧无限迭代

首先编写一个类来演示基本的迭代器协议。这里使用的示例可能与你在其他迭代器教程中看到的示例看上去有所不同，但不要急，因为我认为这种方式能更好地介绍Python中迭代器的工作方式。

接下来的几段内容将实现一个名为Repeater的类，该类可以通过for-in循环迭代，如下所示：

repeater = Repeater('Hello')
for item in repeater:
    print(item)

顾名思义，Repeater类在迭代时类的实例会重复返回同一个值。因此上面的示例代码会一直向控制台输出字符串'Hello'。

为了进行实现，首先定义并填充Repeater类：

class Repeater:
    def __init__(self, value):
        self.value = value

    def __iter__(self):
        return RepeaterIterator(self)

Repeater乍一看像一个普通的Python类，但注意其中包含的__iter__双下划线方法。

__iter__创建并返回了RepeaterIterator对象，这是为了实现for-in迭代功能而必须定义的辅助类：

class RepeaterIterator:
    def __init__(self, source):
        self.source = source

    def __next__(self):
        return self.source.value

同样，RepeaterIterator看起来像一个简单的Python类，但需要注意以下两点。

(1) 在__init__方法中，每个RepeaterIterator实例都链接到创建它的Repeater对象。这样可以持有迭代的“源”（source）对象。

(2) 在RepeaterIterator.__next__中，回到“源”Repeater实例并返回与其关联的值。

在这个代码示例中，Repeater和RepeaterIterator协同工作来支持Python的迭代器协议，其中定义的两个双下划线方法__iter__和__next__是让Python对象迭代的关键。

下面将仔细研究这两个方法，通过对前面介绍的代码进行一些实验来了解其中的工作方式。

首先来确认这两个类的确能让Repeater对象使用for-in循环迭代。为此，先创建一个Repeater实例，迭代时该实例将一直返回字符串'Hello'：

>>> repeater = Repeater('Hello')

现在尝试用for-in循环遍历这个repeater对象。运行以下代码时会发生什么？

>>> for item in repeater:
...     print(item)

不错，屏幕上会显示很多'Hello'。Repeater不断返回相同的字符串值，因此这个循环永远不会停止，会一直向控制台打印'Hello'：

Hello
Hello
Hello
Hello
Hello
...

不过还是恭喜你用Python编写了一个可以工作的迭代器，并用到for-in循环中。虽然循环停不下来，但还算不错。

接下来将剖析这个示例，了解__iter__和__next__方法是如何协同工作来让Python对象迭代的。

有益的提示：如果你在Python REPL会话或终端中运行了上面的示例并且想要停止，请按几次Ctrl + C来跳出无限循环。

Python 迭代器技巧 `for-in`循环在Python中的工作原理

现在已经有了支持迭代器协议的Repeater类，并刚刚运行了一个for-in循环进行了验证：

repeater = Repeater('Hello')
for item in repeater:
    print(item)

那么这个for-in循环在背后究竟做了什么？它如何与repeater对象通信以从中获取新元素？

为了更清楚地说明问题，来将循环展开成一段稍长但结果相同的代码：

repeater = Repeater('Hello')
iterator = repeater.__iter__()
while True:
    item = iterator.__next__()
    print(item)

从中可以看到，for-in只是简单while循环的语法糖。

首先让repeater对象准备迭代，即调用__iter__方法来返回实际的迭代器对象。
然后循环反复调用迭代器对象的__next__方法，从中获取值。

如果你使用过数据库的游标，那就会熟悉这种概念模型：首先初始化游标并准备读取，然后从中逐个取出数据存入局部变量中。

因为在同一时刻只会有一个元素，所以这种方法很节省内存。虽然这个Repeater类提供的是无限长的元素序列，但迭代起来依然没问题。由于无法创建一个包含无限个元素的列表，无法用Python列表模拟相同的行为，因此迭代器是一个非常强大的概念。

用更抽象的术语来说，迭代器提供了一个通用接口，允许在完全隔离容器内部结构的情况下处理容器的每个元素。

无论是元素列表、字典，还是Repeater类提供的无限序列，或是其他序列类型，对于迭代器来说只是实现细节不同。迭代器能以相同的方式遍历这些对象中的元素。

从上面可以看到，Python中的for-in循环没有什么特别之处，在背后都可以都归结为在正确的时间调用__iter__和__next__方法。

实际上，在Python解释器会话中可以手动“模拟”循环使用迭代器协议的方式：

>>> repeater = Repeater('Hello')
>>> iterator = iter(repeater)
>>> next(iterator)
'Hello'
>>> next(iterator)
'Hello'
>>> next(iterator)
'Hello'
...

手动执行的结果与前面相同，每次调用next()时，迭代器都会再次发出相同的问候语。

顺便说一下，这里趁机将__iter__和__next__调用替换为Python的内置函数iter()和next()。
这些内置函数在内部会调用相同的双下划线方法，为迭代器协议提供一个简洁的封装（facade），让代码变得更漂亮、更易读。

Python也为其他功能提供了封装。例如len(x)调用了x.__len__，iter(x)调用了x.__iter__，next(x)调用了x.__next__。

通常最好使用内置的封装函数，不要直接访问实现协议的双下划线方法，这样会让代码更容易阅读。

Python 迭代器技巧更简单的迭代器类

到目前为止的迭代器示例由两个独立的类Repeater和RepeaterIterator组成，直接对应于Python迭代器协议使用的两个阶段。

首先是调用iter()设置和获取迭代器对象，然后通过next()不断从迭代器中获取值。

大部分情况下，这两步可以放到一个类中，用这种方式实现基于类的迭代器代码较少。

之前的第一个例子中没有这样做，因为分开介绍能理清迭代器协议背后的概念模型。现在既然已经明白了如何用烦琐的方法编写一个基于类的迭代器，那么是时候来简化了。

记得为什么还要使用RepeaterIterator类吗？因为要用到这个类中的__next__方法，以便从迭代器中获取新值。不过在哪里定义__next__并不重要。在迭代器协议中，最重要的是__iter__要返回带有__next__方法的对象。

这就诞生了一个想法：RepeaterIterator不断返回相同的值，且不必跟踪任何内部状态。那么能否直接将__next__方法添加到Repeater类中呢？

这样就可以完全摆脱RepeaterIterator，用一个Python类就能实现一个可迭代的对象。来尝试一下，简化后的迭代器示例如下所示：

class Repeater:
    def __init__(self, value):
        self.value = value

    def __iter__(self):
        return self

    def __next__(self):
        return self.value

从含有两个类的10行代码简化成了只有一个类的7行代码，而简化后的实现仍然支持迭代器协议：

>>> repeater = Repeater('Hello')
>>> for item in repeater:
...    print(item)

Hello
Hello
Hello
...

以这种方式简化前面基于类的迭代器通常没有问题。事实上，大多数Python迭代器教程都是直接以这种方式开始介绍，但我始终认为从一开始就用一个类来解释会隐藏迭代器协议的基本原理，增加了理解的难度。

Python 迭代器技巧不想无限迭代

现在你应该很好地掌握了Python中迭代器的工作原理，但是到目前为止只实现了无限迭代的迭代器。

显然，Python中的迭代器主要不是为了无限重复。事实上，回顾本节开头会发现，我使用了下面这个示例来激发大家的学习兴趣：

numbers = [1, 2, 3]
for n in numbers:
    print(n)

你理所当然地期望这段代码输出数字1、2和3后停止，而不是在终端窗口中看到3在不断刷屏，然后在慌乱中狂按Ctrl + C来终止程序。

所以现在来学习如何编写一个会生成新值，并且最终会停下来的迭代器。一般情况下Python对象也不会在for-in循环中无限迭代。

现在来编写另一个名为BoundedRepeater的迭代器类，这个类与之前的Repeater示例类似，但这个类需要能在重复特定次数后停止。

稍微思考一下应该如何做到这一点。迭代器如何表明已执行完毕，没有元素可供迭代了呢？你也许会想：“可以从__next__方法中返回None。”

这个主意不错，但问题在于如果真的希望有些迭代器返回None该怎么办？

来看看其他Python迭代器是如何解决这个问题的。首先来构建一个简单的容器，即包含几个元素的列表，然后遍历这个列表直到耗尽所有元素，观察最后会发生什么：

>>> my_list = [1, 2, 3]
>>> iterator = iter(my_list)
>>> next(iterator)
1
>>> next(iterator)
2
>>> next(iterator)
3

注意，现在已经消耗了列表中的所有三个可用元素，来观察再次在迭代器上调用next会发生什么情况：

>>> next(iterator)
StopIteration

啊哈！引发了StopIteration异常，表示已经耗尽了迭代器中的所有可用值。

没错，迭代器使用异常来处理控制流。为了表示迭代结束，Python迭代器会简单地抛出内置的StopIteration异常。

如果一直向迭代器请求更多的值，就会不断抛出StopIteration异常，表示没有更多的值可供迭代：

>>> next(iterator)
StopIteration
>>> next(iterator)
StopIteration
...

Python迭代器通常不能“重置”，如果其中的元素已经耗尽，那么每次调用next()时都会引发StopIteration。若想重新迭代，需要使用iter()函数获取一个新的迭代器对象。

对于编写在重复一定次数后能停止迭代的BoundedRepeater类，现在万事俱备：

class BoundedRepeater:
    def __init__(self, value, max_repeats):
        self.value = value
        self.max_repeats = max_repeats
        self.count = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.count >= self.max_repeats:
            raise StopIteration
        self.count += 1
        return self.value

这个实现的结果符合预期，迭代在达到max_repeats参数中定义的次数后停止：

>>> repeater = BoundedRepeater('Hello', 3)
>>> for item in repeater:
        print(item)
Hello
Hello
Hello

如果重写上一个for-in循环的例子，移除一些语法糖，那么展开后最终会得到下面的代码片段：

repeater = BoundedRepeater('Hello', 3)
iterator = iter(repeater)
while True:
    try:
        item = next(iterator)
    except StopIteration:
        break
    print(item)

每次在此循环中调用next()时都会检查StopIteration异常，并在必要时中断while循环。
用3行的for-in循环替换8行的while循环是个不错的改进，能让代码更加易读和维护。这样从另一方面看出了Python迭代器的强大之处。

Python 迭代器技巧兼容性

前面展示的所有代码示例都是用Python 3编写的。当涉及实现基于类的迭代器时，Python 2和Python 3之间存在一个小但重要的区别。

在Python 3中，从迭代器中获取下一个值的方法名为__next__。
在Python 2中，相同的方法名为next（不带下划线）。

如果你正在编写基于类的迭代器并试图同时支持Python 2和Python 3，那么这种命名差异会产生一些麻烦。幸运的是，有一种简单的方法可以解决这个问题。
下面是改进后的InfiniteRepeater类，可同时在Python 2和Python 3上运行：

class InfiniteRepeater(object):
    def __init__(self, value):
        self.value = value

    def __iter__(self):
        return self

    def __next__(self):
        return self.value

    # Python 2兼容性：
    def next(self):
        return self.__next__()

为了使这个迭代器类与Python 2兼容，我做了下面两处小改动。

首先添加了一个next方法，简单地调用并返回原__next__的结果。基本上就是为现有的__next__实现创建一个别名以便让Python 2找到。这样就可以支持两个版本的Python，同时所有实际的实现细节仍然位于一个函数中。

其次，将类定义修改为从object继承，以确保在Python 2上创建的是新式类。这与迭代器没有关系，不过是一个很好的习惯。

关键要点