分布式系统,程序语言,算法设计

Python3 生成器(Generator)概念浅析

引子

某次面试问候选人: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 相关的概念主要有三个,分别是 generatorgenerator expressiongenerator 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
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
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