技术赋能#1：Python的可迭代对象、迭代器与生成器

相信大家在学习 Python 的 for 循环的时候，一定或多或少听说过迭代这个词语。例如下面这段代码中的循环可以描述为“变量i对range(100)迭代”，其中i是迭代变量，range(100)是被迭代的对象。而在 Python 中，只有 可迭代对象（iterable） 能被迭代。本文将简要介绍可迭代对象，以及 Python 中与之相关的两个概念：迭代器（iterator）与生成器（generator）。

s = 0
for i in range(100):
    s += i
print(s)

可迭代对象与迭代器

可迭代对象（iterable） 包括两种对象，一是实现了__iter__()方法的对象，其中__iter__()方法要求返回一个迭代器；二是具有序列属性的对象，即具有__getitem__()方法，并可以使用连续的自然数作为索引从对象中获得值。我们所接触的大部分 Python 内置的数据结构的实例都同时满足这两个条件，包括列表list，元组tuple，字符串str和bytes等；而集合set与字典dict只满足第一个条件，不过它们都是可迭代对象。除此之外，range、map、filter和enumerate等内置函数的返回值也是可迭代对象。

迭代器（iterator） 是实现了__next__()方法与__iter__()方法的对象，其中__iter__()方法只能返回对象自身，因此迭代器都是可迭代对象。每次调用该对象的__next__()方法或者对其使用内置函数next()可以从迭代器中获得一个值，当迭代器耗尽时抛出StopIteration异常，结束迭代。

几个示例

上面的描述有点抽象了，还是来看看示例吧。下面的例子以内置函数range为例，演示了可迭代对象和迭代器的一些属性：

>>> rg = range(5)  # 创建range对象
>>> dir(rg)        # rg有__iter__()方法，是可迭代对象
['__bool__',  ..., '__iter__', '__le__', ...]
# 注：next(x)效果同调用 x.__next__()
>>> next(rg)       # rg没有__next__()方法，不是迭代器
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'range' object is not an iterator
# 注：iter(x)效果同调用 x.__iter__()
>>> r = iter(rg)   # 调用__iter__()方法获得可迭代对象对应的迭代器
>>> dir(r)         # r有__iter__()与__next__()方法，是迭代器
['__class__', ..., '__iter__', ..., '__next__', ...]
>>> next(r)        # 使用next()函数或者__next__()方法从r中获得值
0
>>> next(r)
1
>>> r.__next__()
2
>>> r.__next__()
3
>>> next(r)
4
>>> next(r)      # r被耗尽，抛出 StopIteration 异常
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

下面是一个对比示例，演示了 不是迭代器的可迭代对象 和 迭代器 在 for 循环中的区别。

不是迭代器的可迭代对象：代码与执行结果如下。因为每次循环开始前都会调用__iter__()方法返回一个新的迭代器，所以两次循环都是从 0 开始的。

rg = range(5)
for i in rg:
    print(i, end=" ")
    if i>=2: break
print("|", end=" ")
for i in rg:
    print(i, end=" ")
print()
# === output ===
0 1 2 | 0 1 2 3 4

迭代器：将上例中第一行改为rg = iter(range(5))，此时rg变为迭代器，变更后代码的执行结果如下。两次 for 循环开始前都调用__iter__()方法返回rg对象自身（是同一个迭代器），而第一次循环执行到 2 停止了，此时rg中的 0 1 2 均被消耗，因此第二次循环从 3 开始。

0 1 2 | 3 4

自定义迭代器

根据定义，我们知道迭代器需要具有__next__()方法和返回自身的__iter__()方法。不难据此写出自己的迭代器类，下面就是笔者自定义的一个迭代器类，实现了一个可顺序访问的 Fibonacci 数列。

class Fibonacci:
    def __init__(self, count: int):
        self.count = count
        self.last = 1
        self.current = 0

    def __iter__(self):
        # __iter__()方法返回对象自身
        return self

    def __next__(self):
        if self.count == 0: # 如果迭代了指定的次数则抛出 StopIteration
            raise StopIteration()
        # 更新状态
        self.last, self.current = self.current, self.last+self.current
        self.count -= 1
        return self.current # 返回当前值

测试代码：

for i in Fibonacci(15):
    print(i, end=" ")
print()
# === output ===
1 1 2 3 5 8 13 21 34 55 89 144 233 377 610

同时 Python 中的迭代器可以是无限的，在上例中只要将Fibonacci()的参数改为负数即可得到一个无限的迭代器。因为它们是无限的，所以在 for 循环中使用它们时一定不要忘记添加退出条件。

for i in Fibonacci(-1):
    if i > 10000:
        break
    print(i, end=" ")
print()
# === output ===
1 1 2 3 5 8 13 21 34 55 89 144 233 377 610 987 1597 2584 4181 6765

生成器

生成器（生成器函数，generator） 是指 Python 中含有yield关键字的函数，是 Python 提供的用于定义生成器的简易方法。生成器函数在调用时会返回一个生成器迭代器（generator iterator）（顾名思义这是一种迭代器）。

你可能也会觉得上面的那个自定义迭代器的代码太冗长了，我们可以使用生成器将其简化（这里就不考虑无限的情况了）：

def fibonacci(count: int):
    last, current = 1, 0
    for i in range(count):
        last, current = current, last+current
        yield current # 使用yield函数返回当前值，同时表明这是一个生成器

在生成器函数中，当函数执行到yield语句的时候会保存函数当前的状态，然后返回语句里的值，相当于函数的运行在这里暂停了。当__next__()再次被调用的时候函数会加载之前保存的状态，并从yield语句继续执行。你可以在上面的例子中添加一些print函数，结合next对生成器迭代来查看具体的执行过程。

测试代码和结果都和上面的相似，这里就不展示了。需要注意的是，这个例子定义的函数是生成器函数，生成器函数本质上是一个函数，不是可迭代对象，调用生成器函数得到的对象才是迭代器，如下面的例子所示。

>>> fibonacci       # fibonacci是生成器函数
<function fibonacci at 0x7facc4139cf0>
>>> iter(fibonacci) # fibonacci不可迭代
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'function' object is not iterable
>>>
>>> fib = fibonacci(100)
>>> fib             # fibonacci函数调用返回的才是生成器迭代器
<generator object fibonacci at 0x7facc42525e0>
>>> iter(fib)       # fib是可迭代对象，对应的迭代器是其自身（内存地址相同）
<generator object fibonacci at 0x7facc42525e0>
>>> next(fib)
1

生成器推导式

生成器推导式（generator expression） 是对生成器的简化，可以用来替换一些简单的生成器函数。它的语法结构如下，其中<var>为循环变量，<iterable>为被迭代对象，<expr>为输出表达式，<condition>为逻辑表达式（或者是能转换为布尔值的对象）。这个结构与列表推导式和集合推导式相似，除了两边使用的是圆括号，以及返回值是生成器迭代器。

(<expr> for <var> in <iterable> if <condition>)

上面的生成器推导式的结果等价于这个生成器函数的执行结果：

def _():
    for <var> in <iterable>:
        if <condition>:
            yield <expr>

同时也等价于这个 filter-map 链：

map(lambda <var>:<expr>, filter(lambda <var>:<condition>, <iterable>))

关于 yield 关键字

如果你需要在一个生成器中逐个返回另一个可迭代对象的值，比起使用 for 循环+yield，可以用更简便的yield from，比如下面这段代码：

def rangerange(n: int):
    for i in range(n+1):
        yield from range(i)

print(*rangerange(7))
# === output ===
0 0 1 0 1 2 0 1 2 3 0 1 2 3 4 0 1 2 3 4 5 0 1 2 3 4 5 6

还有一个不太常用的知识：yield语句是有返回值的，其值来自于调用生成器迭代器的send()方法时的参数。编程时可以利用这个在迭代过程中从外部影响迭代器的状态。可以用下面这个测试函数看看效果：

def test():
    counter = 0
    while True:
        value = yield counter
        print(f'value received at {counter}:', value)
        counter += 1

# === testing ===
>>> g = test()
>>> next(g)
0
>>> next(g)     # yield语句默认的返回值为None
value received at 0: None
1
>>> g.send(123) # 调用send方法也会触发一次生成器迭代器的更新
value received at 1: 123
2
>>> g.send()    # send方法必须接受一个参数
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: generator.send() takes exactly one argument (0 given)
>>> next(g)
value received at 2: None
3

应用

你可能会想：自定义迭代器到底有什么用呢？我把这些代码移到 for 循环里面不是也有相同的效果？为什么不直接返回一个列表呢？下文将解答你的问题。迭代器的应用与优点包括但不限于：

1. 迭代器是 Python 提供的用于遍历自定义数据结构的接口：假如有人用 Python 实现了一个单向链表，并且希望 for 语句能在这个单向链表上迭代，这个时候就需要在单向链表类中实现__iter__()方法将其变成可迭代对象。
2. 节约内存：以上面的fibonacci生成器为例，记一次性返回完整列表的版本为fibonacci2，假设输入参数是 n，则前者的空间复杂度是 O(1)（因为不用存储每次迭代得到的值），而后者的是 O(n)（需要将每个值存到列表内一次性返回），在 n 很大的时候这个差距尤其明显。
3. 代码封装：针对第二个问题，只有一个 for 循环时确实有相同的效果。若有多个 for 循环都要用到这段代码，则需要对其进行封装。如果使用返回列表的方法，在数据量大的时候会导致内存的浪费，这时生成器无疑是最合适的。

参考资料

• 术语对照表 - Python3.10 官方中文文档