一文掌握 Python 迭代器的原理

理解迭代器是每个严肃的 Python 使用者学习 Python 道路上的里程碑。本文将从零开始,一步一步带你认识 Python 中基于类的迭代器。

相较于其他编程语言,Python 的语法是优美而清晰的。比如 for-in 循环这种极具 Python 特色的代码,让你读起来感觉就像读一个英语句子一样。它很能体现 Python 的美感。

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

但是你知道这种优雅的循环结构背后是如何工作的吗?循环是如何从它所循环的那个对象中获取独立元素的?你怎么才能在自己创建的 Python 对象上使用相同的编程风格呢?

Python 迭代器协议为上边这些疑问提供了答案:

Objects that support the __iter__ and __next__ dunder methods automatically work with for-in loops.

支持 __iter__ 和 __next__ 魔法方法的对象可自动适用于 for-in 循环。

就像 Python 中的装饰器一样,迭代器及其相关技术乍看起来相当神秘而复杂。不用担心,让我们一步一步慢慢来认识它。

我们先聚焦于 Python 3 中迭代器的核心机制,去除不必要的麻烦枝节,这样就可以在基础层面清楚地看到迭代器是如何运作的。

下文中,我们会编写几个支持迭代器协议的 Python 类。每个例子都和上边提到的几个 for-in 循环问题相关联。

好,我们现在进入正题!


【永久迭代的 Python 迭代器】

我们来编写一个可以演示 Python 迭代器协议轮廓的类。你可能学习过其他迭代器教程,但这里使用的例子和你在那些教程中看到的例子有些不同,或许会更有助于你的理解。

我们将会实现一个名为 Repeater 的类,这个类的对象可使用 for-in 循环来迭代。

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

(代码片段1)

诚如类的名字,当被迭代时,该类的对象实例会重复返回一个单一的值。上边这段示例代码会不停地在控制台窗口中打印 Hello 字符串。

我们先来定义这个 Repeater 类:

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

这个类看起来和普通的 Python 类没什么区别。但是请注意它包含了一个 __iter__ 魔法方法。这个方法返回了一个 RepeaterIterator 对象。

这个 RepeaterIterator 对象是什么?它是一个我们需要定义的助手类,借助于 RepeaterIterator,代码片段一中的 repeater 对象才能在 for-in 循环中运行起来。

RepeaterIterator 类的定义如下:

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

同样,它看起来也是一个简单的 Python 类。不过你需要留意以下两点:

  1. 在 __init__ 方法中,我们将每个 RepeaterIterator 实例链接到创建它的 Repeater 对象上。这样,我们就可以保存这个被迭代的 source 对象。
  2. 在 __next__ 方法中,我们又返回了和这个 source Repeater 对象关联的 value 值。

Repeater 和 RepeaterIterator 两个类共同实现了 Python 的迭代器协议。其中的关键是 __iter__ 和 __next__ 这两个魔法方法,正是它们才使得一个 Python 对象成为一个可迭代对象(iterable)。

有了这两个类的定义,你现在可以运行代码片段一了。结果就是,屏幕上不停地打印 Hello 字符串:

Hello
Hello
Hello
Hello
Hello
...

恭喜!你已经实现了一个可以工作的 Python 迭代器,并在 for-in 循环中使用了它。

接下来,我们将拆解这个例子,从而更好地理解 __iter__ 和 __next__ 是如何共同使得一个 Python 对象成为可迭代的。


【for-in 循环是如何工作的】

我们看到,for-in 循环可以从 Repeater 对象中获取新元素,那么,for-in 是如何与 Repeater 对象通信的呢?

我们对代码片段一做些展开,并保持同样的运行结果:

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

如你所见,for-in 就是一个简单 while 循环的语法糖。

  1. 它调用 repeater 对象的 __iter__ 方法获取一个真正的 iterator 对象
  2. 然后在循环中重复调用 iterator 对象的 __next__ 方法,从中获取下一个值

如果你从事过数据库应用开发,使用过数据库游标(cursor),那对此模型应该就不陌生了:先初始化一个游标,将其准备(prepare)好用来读取数据,然后从中获取数据(fetch)并保存到本地变量中。


由于只占用了一个 value 的空间,这种方式具备很高的内存效率。Repeater 类可以用作一个包含元素的无限序列,而这无法通过 Python 列表来实现,我们不可能创建一个包含无限多元素的 list 对象。这是迭代器的一个强大的功能。

使用迭代器的另一个好处是,它提供了一种抽象语义。迭代器可用于多种容器,它为这些容器提供了统一的接口,你不需要关心容器的内部结构就可以从中提取每个元素。

无论你使用列表、字典、无限序列还是别的序列类型,它们都只是代表了一种实现上的细节。而你可以通过迭代器以相同的方式来遍历它们。


我们看到,for-in 循环并无特别之处,它只是在正确的时间调用了正确的魔法方法而已。

实际上,我们也可以在 Python 解释器命令窗口中手动模拟循环是如何使用迭代器协议的。

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

每调用一次 next(),iterator 就会输出一个 Hello。

这里,我们使用 Python 内置的 iter() 和 next() 函数来代替对象的 __iter__ 和 __next__ 方法。这些内置函数其实也是调用了对应的魔法函数,只不过它们为迭代器协议披上了一件干净的外衣,使代码看起来更漂亮更简单。

通常,使用内置函数比直接访问魔法方法要好,因为代码的可读性更强一些。


【一个更简单的迭代器类】

上边的例子中,我们使用两个独立的类来实现迭代器协议,每个类承担协议的一部分职责。而很多时候,这两项职责可以由一个类来承担,从而减少代码量。

我们现在就来简化之前实现迭代器协议的方法。

还记得我们为什么要定义 RepeaterIterator 这个类吗?我们用它来定义获取元素的 __next__ 方法。而实际上,在什么地方定义 __next__ 方法是无关紧要的。

迭代器协议关心的是,__iter__ 方法必须返回一个提供了 __next__ 方法的对象。

我们再审视一下 RepeaterIterator 这个类,它其实并没有做太多的工作,仅仅返回了 source 的成员变量 value。我们是不是可以省去这个类,将其功能融入 Repeater 类中?可以试一下。

我们可以这样来实现新的并且更简单的迭代器类。

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

解释一下:

  1. __iter__ 实现迭代器协议的第一项职责,将 Repeater 自身返回
  2. Repeater 自己定义了 __next__ 方法,每次都返回成员变量 value,从而实现迭代器协议的第二项职责

使用代码片段一来测试这个 Repeater 类,同样会不停地输出 Hello。

通过两个类来实现迭代器协议,我们能很好地理解迭代器协议的底层原则;而使用一个类则可以提升开发效率,非常有意义。


【迭代器不能永远迭代下去】

我们已经理解了迭代器的工作原理,但是我们目前实现的迭代器仅仅可以无限迭代下去,而这并非 Python 中迭代器的主要使用场景。

在本文的开头,我们举了个简单 for-in 循环的例子作为引子:

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

这才是更加普遍的应用场景:输出有限的元素,然后终止。

我们该如何实现这样的迭代器呢?迭代器如何给出元素耗尽迭代结束的信号呢?

或许你会想到,可以在迭代结束时返回 None。听起来是个不错的主意,但是并非适合所有情况,某些场景可能不接受 None 作为合法值。

我们可以看一下 Python 中的其他迭代器是如何处理这一问题的。我们先创建一个简单的容器:一个包含若干元素的 list。然后迭代这个 list,直到元素耗尽,看看会发生什么情况。

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

注意,此时我们已消费了 list 中的所有元素。如果继续调用 next(),会发生什么?

>>> next(iterator)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

报异常了!

没错:迭代器就是使用异常来改变控制流程的。为了提示迭代的结束,Python 迭代器简单地抛出一个内置的 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

BoundedRepeater 可以为我们提供想要的结果。当迭代次数超过 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() 时手动检查是否有异常抛出。for-in 循环替我们做了这项工作,让代码变得更易读和更易维护。这也是 Python 迭代器为何如此强大的另一个原因。


【兼容 Python 2.x 的迭代器】

上文中用到的例子都是使用 Python 3 编写的。如果你使用 Python 2 来实现基于类的迭代器,需要注意一个细小但重要的区别:

  • Python 3 中,从迭代器获取下一个值的方法是:__next__
  • Python 2 中,这个方法叫做:next (没有下划线)

这个命名上的区别可能会导致代码的不兼容。

解决办法也很简单,让迭代器类同时支持这两个方法。

以那个无限序列类为例:

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

    # Python 2 compatibility:
    def next(self):
        return self.__next__()

我们为其增加了一个 next() 方法,该方法仅仅调用 __next__() 就可以了。

还有一个变动,我们将类的定义修改为继承自 object 类,以使其符合 Python 2 中的新式类定义方法。这和迭代器无关,只是一个良好的习惯。


【结语】

  • 迭代器为 Python 对象提供了一个内存效率高的序列接口,使用起来也更具 Python 风格。
  • Python 对象可通过实现 __iter__ 和 __next__ 魔法方法的方式来支持迭代。
  • 基于类的迭代器是 Python 中的一种编写可迭代对象的方法,除此之外,还可以考虑使用生成器和生成器表达式。

【近期热门文章】

  1. 从 Python 列表的特性来探究其底层实现机制

  2. 从 Python 源码来分析列表的 resize 机制

  3. 列表推导式:简洁高效更具 Python 风格的列表创建方法

  4. Python 列表的应用场景有哪些?你使用对了吗?


欢迎转载,请注明出处。谢绝搬运,抄袭必究!

欢迎关注本站公众号【python学与思】

python 学与思