python iterator 遇到的一个坑

在使用 Python 的 iterator 时,遇到一个很愚蠢的错误,浪费的很多时间才找到原因。特此记录一下,提醒自己,提示他人。
Binder

Bug 复现

原来的问题比较复杂,导致出错难以调试,经过将问题不断简化,最后简化后的代码如下:

打印 iterator 的函数

我们定义一个用于打印 iterator 的函数

1
2
3
4
5
6
7
8
def print_iterator(iterator):
while True:
try:
element = next(iterator)
except StopIteration:
break
else:
print(element)

测试一下

1
print_iterator(iter([1, 2, 3]))

输出:

1
2
3
1
2
3

代码工作正常

定义一个计算 iterable 长度的函数

1
2
3
4
5
def counter_iterable(iterable):
iterator = iter(iterable)
iterator_length = sum(1 for _ in iterator)

print(iterator_length)

测试一下

1
counter_iterable([1, 2, 3])

输出:

1
3

代码工作正常

整合在一起

放在 中调用:

1
2
3
4
5
6
7
def counter_iterable_and_print(iterable):
iterator = iter(iterable)
iterator_length = sum(1 for _ in iterator)

print(iterator_length)

print_iterator(iterator)

测试一下

1
counter_iterable_and_print([1, 2, 3])

期望的输出应该有长度和打印内容两个部分构成:

1
2
3
4
3
1
2
3

但实际上输出是:

1
3

代码工作不正常

bug 原因

这个 bug 的产生和我对两个概念的理解和记忆错误有关:
第一个是没有充分理解和记忆 iterator 的工作机制。

iterator 是什么

根据 Iterator on Python wiki :

An iterator is an object that implements next (in python3, it is __next__) method, which is expected to return the next element of the iterable object that returned it, and raise a StopIteration exception when no more elements are available.

Iterator will typically need to maintain some kind of position state information (like the index of the last element returned or the like). If the iterable maintained that state itself, it would become inherently non-reentrant (meaning you could use it only one loop at a time).

上述 bug 产生的原因:上面的代码中共享了一个 iterator 对象,由于 iterator 具有记忆内部状态的能力,所以当

1
iterator_length = sum(1 for _ in iterator)

执行完毕后,实际这个 iterator 对象已经完成了全部元素的迭代。后续再次调用这个对象的 __next__() 方法时,直接抛出 StopIteration 异常。因此这就是为什么后续的

1
print_iterator(iterator)

并没有任何输出的原因,因为函数收到的参数已经是一个走到最后的 iterator 对象了。

iterator 内有有状态信息,具有不可重入(non-reentrant)的特性,这个和 list, tuple, dict 等容器不一样,容器通过 __getitem__ 来迭代。

另外一个是没有搞清楚 iterableiterator 的区别。正确的理解是:iterable 是一个工厂函数,通过显式的调用 iter 函数或者调用其 __iter__() 方法或者使用 for 循环来获得这个工厂的产品:一个 iterator 对象。

这里需要简单说明的是:无论是使用 iter 函数还是在 for 循环中使用,都是间接的调用 iterable 对象的 __iter__ 方法。

iterator 被要求需要支持 iterable 协议的。通常情况下,对 iterator 调用 __iter__ 方法,返回的 iterator 就是它自己。 这样 iterator 对象就能够在 for-loop 中使用了。

关于 iterator 的 PEP 在 PEP 234 – Iterators