Python3 生成器（generator）概念浅析

「这是我参与2022首次更文挑战的第 27 天，活动详情查看：2022首次更文挑战」。

引子

某次面试问候选人：Python 中生成器是什么？答曰：有 yield 关键字的函数。而在我印象中此种函数返回的值是生成器，而函数本身不是。如下：

In [1]: def get_nums(n): 
   ...:     for i in range(n): 
   ...:         yield i 
   ...:                                                                                                                                                                  
In [2]: type(get_nums)
Out[2]: function
  
In [3]: nums = get_nums(10)   
  
In [4]: type(nums)
Out[4]: generator
复制代码

但看候选人那么笃定，隐隐然感觉哪里不对，于是有了以下探究。

作者：木鸟杂记 www.qtmuniao.com, 转载请注明出处

概念

要弄清楚 Python3 中这些概念的区别，最权威的当然是去看官方的术语对照表：

generator

A function which returns a generator iterator. It looks like a normal function except that it contains yield expressions for producing a series of values usable in a for-loop or that can be retrieved one at a time with the next() function.

Usually refers to a generator function, but may refer to a generator iterator in some contexts. In cases where the intended meaning isn’t clear, using the full terms avoids ambiguity.

generator iterator

An object created by a generator function.

Each yield temporarily suspends processing, remembering the location execution state (including local variables and pending try-statements). When the generator iterator resumes, it picks up where it left off (in contrast to functions which start fresh on every invocation).

generator expression

An expression that returns an iterator. It looks like a normal expression followed by a for clause defining a loop variable, range, and an optional if clause.

可以看出，Python generator 相关的概念主要有三个，分别是 generator，generator expression，generator iterator。上面也特别指出了， genrator 在有的上下文中指的是 generator function（如候选人所说），但在另外一些上下文中却指的是 generator iterator（如上面例子解释器告诉我们的）。为了避免歧义，大家在表达时可以尽量说全称。

结合我的一些经验，可以将其归纳为以下三个概念：

Generator Function：含有 yield 关键字的函数，会返回一系列值，可以使用 next() 对其返回值进行迭代。
Generator Iterator：generator function 返回的对象。可以进行一次性地迭代。
Generator Expression：可以求值为 generator iterator 的表达式。使用小括号和 for 来定义，如下面例子：

In [5]: a = (i*i for i in range(10))                                                                                                                                                                                                                                     

In [6]: type(a)                                                                                                                                                                                                                                                          
Out[6]: generator
复制代码

深入

生成器 vs 迭代器

从中文字面可能不好理解它们的关系，但是从上文提到的英文术语来分析：generator iterator。他们的关系就一目了然了：

迭代器（Iterator）是一种更宽泛的概念，生成器（generator Iterator）是一种迭代器，但反过来不成立。

迭代器是任何实现了 __next__ 方法的对象（object），可以通过 next(iterator) 对其进行迭代，迭代结束时会抛出 StopIteration 异常。

while True:
  try:
    x = next(an_iterator)
    do_sth_with(x)
  except StopIteration:
    break
 
复制代码

通常我们会使用 for in 来对其进行简化迭代：

for x in an_iterator:
  do_sth_with(x)
复制代码

yield 原理

yield 是一个神奇的关键字，它会临时挂起当前函数，记下其上下文（包括局部变量、待决的 try catch 等），将控制权返回给函数调用者。当下一次再调用其所在 generator function 时，会恢复保存的上下文，继续执行剩下的语句，直到再遇到 yield 或者退出为止。

我们常见的 return 是与之相对的关键字，但 return 会结束函数调用，销毁上下文（弹出栈帧），将控制权返回给调用者。

因此，以 yield 进行执行流控制的函数称为 generator function，以 return 进行执行流控制的函数，就是普通的 function 喽~

当然，由于可以临时挂起函数的执行，yield 还有更高阶的用法，即充当其调用者和被挂起函数间交互的桥梁：

In [1]: def dynamic_step_seq(size, start=0, default_step=1): 
    ...:   x = start 
    ...:   for _ in range(size): 
    ...:     given_step = yield x 
    ...:     if given_step is not None: 
    ...:       x += given_step 
    ...:     else: 
    ...:       x += default_step 
    ...:                                                                                                                                                                                                                                                                 

In [2]: for x in dynamic_step_seq(10): 
    ...:     print(x, end=' ') 
    ...: print()                                                                                                                                                                                                                                                         
0 1 2 3 4 5 6 7 8 9 
In [3]: dss = dynamic_step_seq(10)                                                                                                                                                                                                                                      

In [4]: next(dss) # 注                                                                                                                                                                                                                                                
Out[5]: 0

In [6]: dynamic_step = 1 
    ...: while True: 
    ...:     try: 
    ...:         x = dss.send(dynamic_step) 
    ...:         dynamic_step += 1 
    ...:         print(x, end=' ') 
    ...:     except StopIteration: 
    ...:         print() 
    ...:         break 
    ...:                                                                                                                                                                                                                                                                 
1 3 6 10 15 21 28 36 45 
  
复制代码

注：此处初看有些奇怪，但是通过 yield 作用我们能推断出原理：需要首先调用 next 将函数运行至 yield 处，才能通过 generator.send 给 generator 传送对象。

效用

那么使用生成器有什么好处呢？简单来说，主要有两大好处：

精简代码
提高性能

精简代码

使用 yield 关键字或者生成器表达式可以很方便的生成一个迭代器对象。为了说明这一点，我们来比较一下对于一个需求的不同实现。该需求很简单：获取前 n 个自然数。

最直观的方法为，构造一个数组然后返回：

# Build and return a list
def firstn(n):
    num, nums = 0, []
    while num < n:
        nums.append(num)
        num += 1
    return nums

sum_of_first_n = sum(firstn(1000000))
复制代码

当 n 很小的时候，该实现没有什么问题，但是当 n 变得很大，你的机器内存是吃不消的：

In [4]: a = firstn(10000000000000000000)                                                                                                                                                                                                                                 
Killed
复制代码

IPython 直接内存爆掉被 kill 了。

于是，我们很自然的想起可以用生成器模式，但是你仍然不想用 yield，于是你需要构造一个对象，并实现 __iter__ 和 __next__ 方法：

# Using the generator pattern (an iterable)
class firstn(object):
    def __init__(self, n):
        self.n = n
        self.num, self.nums = 0, []

    def __iter__(self):
        return self

    def __next__(self):
        if self.num < self.n:
            cur, self.num = self.num, self.num+1
            return cur
        else:
            raise StopIteration()

sum_of_first_n = sum(firstn(1000000))
复制代码

注：在 Python3 中，如果你的类实现了__iter__ ，则成为一个可迭代对象，可以调用 iter(instance) 得到一个迭代器。如果你的类实现了 __next__，那么你的类实例本身就成为了一个迭代器，可以通过 next(instance) 来调用，进行迭代。

这些代码终于实现了我们的要求。但是，它有如下问题：

为了实现一个简单的需求却不得不构造冗长的代码。
使用了很多语言约定，读起来不直观。
逻辑表达的绕来绕去。

于是，你忍不住了，终于使用了 yield：

# a generator that yields items instead of returning a list
def firstn(n):
    num = 0
    while num < n:
        yield num
        num += 1

sum_of_first_n = sum(firstn(1000000))
复制代码

于是你使用了寥寥几行，构造出了一个生成器函数。

其实[坏笑]，Python 3 中，就有该函数： range(n)。

利用此函数以及 Iterator expression，你可以用很紧凑的代码构造出很多迭代数列：

squares = (x * x for x in range(n))
evens = (2 * x for x in range(n))
复制代码

注意，他们都只能迭代一次。

提高性能

这一条主要是针对内存使用上来说的。因为迭代器不会保存所有值，而是在运行中动态的计算出数列的各个值，并将之前的数值扔掉，因此是很省内存的，这点在前面的例子也有体现。这里举另外一个更实际一点例子。

假设我们有一个很大的文件（比如说 8G），但是你的电脑只有 4G 内存，你如何利用 Python 对其进行处理？

答案是使用 yield 构造 generator Iterator：

def read_by_chunks(file, chunk_size=1024):
    while True:
        data = file.read(chunk_size)
        if not data:
            break
        yield data


f = open('your_big_file.dat')
for chunk in read_by_chunks(f):
    process_chunk(chunk)
复制代码

这种通过构造 generator 逐块读取的方法又叫惰性加载，也叫流式读取，是处理大文件的一种常见方式。如果你的是文本文件，可以按行读取，那代码就更简单了：

with open('your_big_file.txt') as f:
    for line in f: 
        process_line(line)
复制代码

Python3 文件类句柄就是一个迭代器，默认会将文件按行分割以惰性加载。

参考

Python 3 术语对照表：docs.python.org/3/glossary.…
Python 3 wiki：wiki.python.org/moin/Genera…
Iterator vs Iterable: stackoverflow.com/questions/5…
Lazy read big file：stackoverflow.com/questions/5…

我是青藤木鸟，一个喜欢摄影的分布式系统程序员，欢迎关注我的公众号：“木鸟杂记”