Python 生成器 ~ 看一篇就够用

一. 生成器表达式与列表推导

小编的另外一篇文章 Python -> 每日 GET 一个新技能 之 Python 列表推导，介绍了列表推导相较于 map filter 的种种好处。

但是，在输入的数据量较大时，列表推导很可能会由于占用大量内存而导致 memory leak，特别是使用网络套接字来无休止地读取时，那么列表推导就会出问题。

这时，我们就需要用生成器表达式来改写数据量较大的列表推导。由于生成器表达式在运行时，并不会把整个输出序列都呈现出来，而是会估值为迭代器（iterator），这个迭代器每次可以根据生成器表达式产生一项数据。

下面是一组对比演示：

Python 生成器 ~ 看一篇就够用

>>> with open( test_generator.txt ) as f:
           line_lens = [len(x) for x in f]
>>> line_lens
[2, 3, 4, 5, 6, 7, 8, 9, 10]

把实现列表推导的写法放在一堆圆括号中，就构成了生成器表达式：

>>> line_it = (len(x) for x in open( test_generator.txt ))
>>> line_it
<generator object <genexpr> at 0x000001CA972F1C10>

把此生成器作为参数，传递给 next 函数，即可逐次产生输出值，从而避免了内存用量问题：

>>> next(line_it)
2
>>> next(line_it)
3

此外，把某个生成器表达式所返回的迭代器，放在另一个生成器表达式的 for 子表达式中，即可将二者串起来：

>>> it2 = ((x, x**0.5) for x in line_it)
>>> it2
<generator object <genexpr> at 0x000001CA97354190>
>>> next(it2)
(4, 2.0)
>>> next(it2)
(5, 2.23606797749979)

另外串在一起的生成器表达式执行速度很快。

注意：由生成器表达式所返回的那个迭代器是有状态的，用过一轮之后，请勿反复使用。细心的同学也会发现，it2 是在 line_it 被 next 过 2 次之后才用于 it2 ，故 it2 第一次 next 时，就是读取的 aaa ，返回 (4, 2.0)

当所有的行都被迭代完之后，会抛出 StopIteration 异常：

>>> for i in it2:
           print(i)
(6, 2.449489742783178)
(7, 2.6457513110645907)
(8, 2.8284271247461903)
(9, 3.0)
(10, 3.1622776601683795)
>>> next(it2)
---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-77-58c7b9b7f185> in <module>
----> 1 next(it2)

StopIteration:

二. `iter` 方法实现为生成器，即可定义自己的容器类

通过上面的学习，我们已经了解到迭代器是有状态的，用过一轮之后，会抛出 StopInteration 异常，后续无法重复使用。下面，我们再来看一个演示：
函数 print_precent 计算参数 numbers 中每一项所占的 sum(numbers) 比重：

def print_percent(numbers):
    total = sum(numbers)
    for value in numbers:
        percent = 100 * value / total
        print("%d percent: %.2f%%" % (value, percent))

调用时，传入容器类型 list ，结果完全符合我们的预期：

>>> print_percent([15, 35, 80])
15 percent: 11.54%
35 percent: 26.92%
80 percent: 61.54%

下面我们定义一个生成器函数，返回文件中各行的长度：

def read_numbers(data_path):
    with open(data_path) as f:
        for line in f:
            yield len(line)

包含 yield 表达式的函数即为生成器函数，调用生成器函数，返回迭代器：

>>> read_numbers("test_generator.txt")
<generator object read_numbers at 0x000001CA973CBAC0>

下面我们将次生成器函数的调用结果传递给 print_percent：

>>> it = read_numbers("test_generator.txt")
>>> print_percent(it)
>>>

这里，奇怪的事情便发生了：以生成器函数所返回的迭代器作为参数，来调用 print_percent ，没有产生任何打印。下面我们来一探究竟：

>>> it = read_numbers("test_generator.txt")
>>> list(it)
[2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> list(it)
[]

>>> it = read_numbers("test_generator.txt")
>>> sum(it)
54
>>> sum(it)
0

没错，细心的同学可能已经发现了：list 和 sum 在使用迭代器时，都会进行迭代。而迭代器只能产生一轮结果，在抛出 StopIteration 异常的迭代器或生成器上面继续迭代第二轮，是不会有结果的。

注：for 循环、list 构造器以及 sum 等函数，都认为在正常操作过程中完全可能出现 StopIteration 异常，他们没办法区分这个迭代器是本来就没有值，还是本来有值，但已经用完了。因此，在已用完的迭代器上面继续迭代时，是不会报错的。

为解决 print_percent 参数迭代的问题，下面我们尝试三种不同方向的解决办法。三种讨论之后，信任你定然会找到最优解。

尝试一

使用迭代器制作一份列表，将它的全部内容都遍历一次，并复制到列表中：

def print_percent_copy(numbers):
    numbers = list(numbers)
    total = sum(numbers)
    for value in numbers:
        percent = 100 * value / total
        print("%d percent: %.2f%%" % (value, percent))

运行结果：

>>> it = read_numbers("test_generator.txt")
>>> print_percent_copy(it)
2 percent: 3.70%
3 percent: 5.56%
4 percent: 7.41%
5 percent: 9.26%
6 percent: 11.11%
7 percent: 12.96%
8 percent: 14.81%
9 percent: 16.67%
10 percent: 18.52%
>>>

调用改善后的 print_percent_copy 函数，能够正确打印结果了。但这种方法的问题便在于大数据量了，如果迭代器包含大量输入数据，很可能会导致程序在复制迭代器的时候耗尽内存并崩溃。

尝试二

既然不能将函数的参数轻易的进行复制，那么另一种解决办法，就是通过传入一个函数，该函数每次调用后都会返回一个新的迭代器：

def print_percent_callabel_arg(get_iter):
    total = sum(get_iter())
    for value in get_iter():
        percent = 100 * value / total
        print("%d percent: %.2f%%" % (value, percent))

运行结果：

>>> print_percent_callabel_arg(lambda: read_numbers("test_generator.txt"))
2 percent: 3.70%
3 percent: 5.56%
4 percent: 7.41%
5 percent: 9.26%
6 percent: 11.11%
7 percent: 12.96%
8 percent: 14.81%
9 percent: 16.67%
10 percent: 18.52%

上述的 lambda 表达式被调用时，返回生成器，以便每次都能产生新的迭代器。

尝试三

尝试二 虽然得到了较好的解决，但看上去并不那么流畅，略显复杂。实则，在 Python 中利用迭代器协议，是可以实现自己的容器类的。

注：Python 在 for 循环及相关表达式中遍历某容器的内容时，依赖的正是迭代器协议。举例来讲：在迭代 o 对象时，Python 第一会调用 iter(o) ，内置的 iter 函数又会去调用特殊方法 o.__iter__，而 __iter__必须返回迭代器。
我们知道迭代器本身实现了 __next__ 特殊方法，之后，for 循环会在迭代器上面反复调用 next 函数，直至耗尽产生 StopIteration 异常。

__iter__ 方法既然必须返回迭代器，利用这一特点，我们刚好把 __iter__ 实现为我们上述的生成器函数，即可定义出自己的容器类，随意多少次迭代，随意输入数据有多大，都可以满足要求。

def print_percent(numbers):
    total = sum(numbers)
    for value in numbers:
        percent = 100 * value / total
        print("%d percent: %.2f%%" % (value, percent))


class ReadNumbers:
    def __init__(self, data_path):
        self.data_path = data_path
        
    def __iter__(self):
        with open(self.data_path) as f:
            for line in f:
                yield len(line)

print_percent 函数中的 sum 方法，会调用 ReadNumbers.__iter__，得到一个新的迭代器，后面紧跟的 for 循环，也会调用 __iter__，从而得到另外一个新的迭代器。

运行结果：

>>> numbers = ReadNumbers("test_generator.txt")
>>> print_percent(numbers)
2 percent: 3.70%
3 percent: 5.56%
4 percent: 7.41%
5 percent: 9.26%
6 percent: 11.11%
7 percent: 12.96%
8 percent: 14.81%
9 percent: 16.67%
10 percent: 18.52%

清楚了容器的工作原理之后，最后我们还需要改善 print_percent 函数，以确保调用者传递的参数，并不是迭代器对象本身。

关于迭代器协议有这样的约定：如果把迭代器对象传递给内置的 iter 函数，那么此函数会把该迭代器返回；若传给 iter 函数的是个容器类对象，那么 iter 函数每次都会返回新的迭代器对象。

Python 生成器 ~ 看一篇就够用

利用这个特点，可以在 print_percent 中添加如下判断：

def print_percent(numbers):
    if iter(numbers) is iter(numbers):
        raise TypeError( Must supply a container )
    total = sum(numbers)
    for value in numbers:
        percent = 100 * value / total
        print("%d percent: %.2f%%" % (value, percent))

改善后的 print_percent 函数，在传入我们前面定义的 ReadNumbers 容器类实例时，可以正常运行：

>>> print_percent(numbers)
2 percent: 3.70%
3 percent: 5.56%
4 percent: 7.41%
5 percent: 9.26%
6 percent: 11.11%
7 percent: 12.96%
8 percent: 14.81%
9 percent: 16.67%
10 percent: 18.52%

下面，我们为改善后的 print_percent 函数，传入迭代器对象，则会抛出 TypeError 异常，以此来确保输入的参数类型为容器类，直接传入迭代器，则会抛出异常。

>>> it = read_numbers("test_generator.txt")
>>> it
<generator object read_numbers at 0x00000286922D4900>
>>> print_percent(it)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-12-86e70d62a08e> in <module>
----> 1 print_percent(it)

<ipython-input-7-ddb2e7adb29f> in print_percent(numbers)
      1 def print_percent(numbers):
      2     if iter(numbers) is iter(numbers):
----> 3         raise TypeError( Must supply a container )
      4     total = sum(numbers)
      5     for value in numbers:

TypeError: Must supply a container

总结：
① 函数对输入参数多次迭代需要注意：如果参数是迭代器，那么后续的迭代将会是无效的；
② Python 的迭代器协议，描述了容器和迭代器应该如何与 iter 和 next 内置函数、for 循环及相关表达式相互配合；
③ 把 __iter__ 方法实现为生成器，即可定义自己的容器类型；
④ 判断某个值是迭代器还是容器，可以拿该值为参数，两次调用 iter 函数，若结果一样，则是迭代器；相反，则为容器。