木鸟杂记

分布式系统，数据库，存储

Python3 生成器（Generator）概念浅析

发表于 更新于 分类于 ，

python-generator.png

引子

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

1
2
3
4
5
6
7
8
9
10
11
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

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

作者：木鸟杂记 https://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 来定义，如下面例子：
1
2
3
4
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 异常。

1
2
3
4
5
6
7
while True:
  try:
    x = next(an_iterator)
    do_sth_with(x)
  except StopIteration:
    break

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

1
2
for x in an_iterator:
  do_sth_with(x)

yield 原理

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

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

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

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
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 个自然数。

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

1
2
3
4
5
6
7
8
9
# 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 变得很大，你的机器内存是吃不消的：

1
2
In [4]: a = firstn(10000000000000000000)                                                                                                                                                                                                                                 
Killed

IPython 直接内存爆掉被 kill 了。

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 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) 来调用，进行迭代。

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

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

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

1
2
3
4
5
6
7
8
# 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，你可以用很紧凑的代码构造出很多迭代数列：

1
2
squares = (x * x for x in range(n))
evens = (2 * x for x in range(n))

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

提高性能

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

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

答案是使用 yield 构造 generator Iterator：

1
2
3
4
5
6
7
8
9
10
11
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 逐块读取的方法又叫惰性加载，也叫流式读取，是处理大文件的一种常见方式。如果你的是文本文件，可以按行读取，那代码就更简单了：

1
2
3
with open('your_big_file.txt') as f:
    for line in f: 
        process_line(line)

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

参考

  1. Python 3 术语对照表：https://docs.python.org/3/glossary.html

  2. Python 3 wiki：https://wiki.python.org/moin/Generators

  3. Iterator vs Iterable: https://stackoverflow.com/questions/52056146/separating-the-iter-and-next-methods/52056290

  4. Lazy read big file：https://stackoverflow.com/questions/519633/lazy-method-for-reading-big-file-in-python

我是青藤木鸟，一个喜欢摄影、专注大规模数据系统的程序员，欢迎关注我的公众号：“木鸟杂记”，有更多的分布式系统、存储和数据库相关的文章，欢迎关注。 关注公众号后，回复“资料”可以获取我总结一份分布式数据库学习资料。 回复“优惠券”可以获取我的大规模数据系统付费专栏《系统日知录》的八折优惠券。

我们还有相关的分布式系统和数据库的群，可以添加我的微信号：qtmuniao，我拉你入群。加我时记得备注：“分布式系统群”。 另外，如果你不想加群，还有一个分布式系统和数据库的论坛（点这里），欢迎来玩耍。

wx-distributed-system-s.jpg