Python生成器详解

生成器和迭代器的功能非常相似,它也会提供 __next__() 方法,这意味着程序同样可调用内置的 next() 函数来获取生成器的下一个值,也可使用 for 循环来遍历生成器。

生成器与迭代器的区别在于,迭代器通常是先定义一个迭代器类,然后通过创建实例来创建迭代器;而生成器则是先定义一个包含 yield 语句的函数,然后通过调用该函数来创建生成器。

生成器是一种非常优秀的语法,Python 使用生成器可以让程序变得很优雅。

创建生成器

创建生成器需要两步操作:
  1. 定义一个包含 yield 语句的函数。
  2. 调用第 1 步创建的函数得到生成器。

下面程序使用生成器来定义一个差值递增的数列。程序先定义了一个包含 yield 语句的函数:
def test(val, step):
    print("--------函数开始执行------")
    cur = 0
    # 遍历0~val
    for i in range(val):
        # cur添加i*step
        cur += i * step
        yield cur
上面函数与前面介绍的普通函数的最大区别在于 yield cur 这行,如果将这行代码改为 print(cur),那么这个函数就显得比较普通了,该函数只是简单地遍历区间,并将循环计数器乘以 step 后添加到 cur 变量上,该数列中两个值之间的差值会逐步递增。

如果将上面的 yield cur 语句改为 print(cur, end =''),执行 test(10, 2) 函数将会看到如下输出结果:

--------函数开始执行------
0 2 6 12 20 30 42 56 72 90

yield cur 语句的作用有两点:
  1. 每次返回一个值,有点类似于 return 语句。
  2. 冻结执行,程序每次执行到 yield 语句时就会被暂停。

在程序被 yield 语句冻结之后,当程序再次调用 next() 函数获取生成器的下一个值时,程序才会继续向下执行。

需要指出的是,调用包含 yield 语句的函数并不会立即执行,它只是返回一个生成器。只有当程序通过 next() 函数调用生成器或遍历生成器时,函数才会真正执行。

保留上面函数中的 yield cur 语句,执行如下语句:
# 执行函数,返回生成器
t = test(10, 2)
print('=================')
# 获取生成器的第一个值
print(next(t)) # 0,生成器“冻结”在yield处
print(next(t)) # 2,生成器再次“冻结”在yield处
运行上面代码,可以看到如下输出结果:

=================
--------函数开始执行------
0
2

从上面的输出结果不难看出,当程序执行 t = test(10, 2) 调用函数时,程序并未开始执行 test() 函数;当程序第一次调用 next(t) 时,test() 函数才开始执行。

Python 2.x 不使用 next() 函数来获取生成器的下一个值,而是直接调用生成器的 next() 方法。也就是说,在 Python 2.x 中应该写成 t.next()。

当程序调用 next(t) 时,生成器会返回 yield cur 语句返回的值(第一次返回 0),程序被“冻结”在 yield 语句处,因此可以看到上面生成器第一次输出的值为 0。

当程序第二次调用 next(t) 时,程序的“冻结”被解除,继续向下执行,这一次循环计数器 i 变成 1,在执行 cur += i * step 之后,cur 变成 2 ,生成器再次返回 yield cur 语句返回的值(这一次返回 2),程序再次被“冻结”在该 yield 语句处,因此可以看到上面生成器第二次输出的值为 2。

程序可使用 for 循环来遍历生成器,相当于不断地使用 next() 函数获取生成器的下一个值。例如如下代码:
for ele in t:
print(ele, end=' ')
运行上面循环代码,会生成如下输出结果:

6 12 20 30 42 56 72 90

由于前面两次调用 next() 函数已经获取了生成器的前两个值,因此此处循环时第一次输出的值就是 6。

此外,程序可使用 list() 函数将生成器能生成的所有值转换成列表,也可使用 tuple() 函数将生成器能生成的所有值转换成元组。例如如下代码:
#再次创建生成器
t = test(10, 1)
#将生成器转换成列表
print (list (t))
#再次创建生成器
t = test(10, 3)
#将生成器转换成元组
print(tuple(t))
运行上面代码,可以看到如下输出结果:

--------函数开始执行------
[0 , 1 , 3, 6 , 10 , 15 , 21 , 28 , 36 , 45]
--------函数开始执行------
(0 , 3 , 9 , 18 , 30 , 45 , 63 , 84 , 108, 135)

系统学习本教程的读者应该知道,前面章节还介绍过使用 for 循环来创建生成器(将 for 表达式放在圆括号里)。可见,Python 主要提供了以下两种方式来创建生成器:
  1. 使用 for 循环的生成器推导式。
  2. 调用带 yield 语句的生成器函数。

生成器是 Python 的一个特色功能,在其他语言中往往没有对应的机制,因此很多 Python 开发者对生成器机制不甚了解。但实际上生成器是一种非常优秀的机制,以我们实际开发的经验来看,使用生成器至少有以下几个优势:
  • 当使用生成器来生成多个数据时,程序是按需获取数据的,它不会一开始就把所有数据都生成出来,而是每次调用 next() 获取下一个数据时,生成器才会执行一次,因此可以减少代码的执行次数。比如前面介绍的示例,程序不会一开始就把生成器函数中的循环都执行完成,而是每次调用 next() 时才执行一次循环体。
  • 当函数需要返回多个数据时,如果不使用生成器,程序就需要使用列表或元组来收集函数返回的多个值,当函数要返回的数据量较大时,这些列表、元组会带来一定的内存开销。如果使用生成器就不存在这个问题,生成器可以按需、逐个返回数据。
  • 使用生成器的代码更加简洁。

生成器的方法

当生成器运行起来之后,开发者还可以为生成器提供值,通过这种方式让生成器与“外部程序”动态地交换数据。

为了实现生成器与“外部程序” 动态地交换数据,需要借助于生成器的 send() 方法,该方法的功能与前面示例中所使用的 next() 函数的功能非常相似,它们都用于获取生成器所生成的下一个值,并将生成器“冻结”在 yield 语句处;但 send() 方法可以接收一个参数,该参数值会被发送给生成器函数。

在生成器函数内部,程序可通过 yield 表达式来获取 send() 方法所发送的值,这意味着此时程序应该使用一个变量来接收 yield 语句的值。如果程序依然使用 next() 函数来获取生成器所生成的下一个值,那么 yield 语句返回 None。

对于上面详细的描述,归纳起来就是两句话:
  1. 外部程序通过 send() 方法发送数据。
  2. 生成器函数使用 yield 语句接收收据。

另外,需要说明的是,只有等到生成器被“冻结”之后,外部程序才能使用 send() 方法向生成器发送数据。获取生成器第一次所生成的值,应该使用 next() 函数;如果程序非要使用 send() 方法获取生成器第一次所生成的值,也不能向生成器发送数据,只能为该方法传入 None 参数。

下面程序示范了向生成器发送数据。该程序会依次生成每个整数的平方值,但外部程序可以向生成器发送数据,当生成器接收到外部数据之后会生成外部数据的平方值:
def square_gen(val):
    i = 0
    out_val = None
    while True:
        # 使用yield语句生成值,使用out_val接收send()方法发送的参数值
        out_val = (yield out_val ** 2) if out_val is not None else (yield i ** 2)
        # 如果程序使用send()方法获取下一个值,out_val会获取send()方法的参数
        if out_val is not None : print("====%d" % out_val)
        i += 1

sg = square_gen(5)
# 第一次调用send()方法获取值,只能传入None作为参数
print(sg.send(None))  # 0
print(next(sg))  # 1
print('--------------')
# 调用send()方法获取生成器的下一个值,参数9会被发送给生成器
print(sg.send(9))  # 81
# 再次调用next()函数获取生成器的下一个值
print(next(sg))  # 9
该程序与前面的简单生成器程序的区别就在于第 6 行代码,这行代码在 yield 语句(yield 语句被放在 if 表达式中,整个表达式只会运回一个 yield 语句)的左边放了一个变量,该变量就用于接受生成器 send() 方法所发送的值。

上面程序第一次使用生成器的 send() 方法来获取生成器的下一个值,因此只能为 send() 方法传入 None 作为参数。程序执行到第 9 行代码,由于此时 out_val 为 None ,因此程序执行 yield i**2(生成器返回 0),程序被“冻结”。注意,当程序被“冻结” 时,程序还未对 out_val 变量赋值,因此看到第一次获取生成器的值为 0 。

通过上面的执行过程不难看出,生成器根本不能获取第一次调用 send() 方法发送的参数值,这就是 Python 要求生成器第一次调用 send() 方法时只能发送 None 参数的原因。

接下来程序调用 next(sg) 获取生成器的下一个值,程序从“冻结”处(对 out_val 赋值)向下执行。由于此处调用 next() 函数获取生成器的下一个值,因此 out_val 被赋值为 None,所以程序执行 yield i**2(生成器返回 1),程序再次被“冻结”。

接下来程序调用 sg.send(9) 获取生成器的下一个值,程序从“冻结”处(对 out_val 赋值)向下执行。由于此处调用 send(9) 方法获取生成器的下一个值,因此 out_val 被赋值为 9,所以程序执行 yield out_val**2(生成器返回 81),程序再次被“冻结”。因此看到本次获取生成器的值为 81。

程序再次调用 next(sg) 获取生成器的下一个值,程序从“冻结”处(对 out_val 赋值)向下执行。由于此处调用 next() 函数获取生成器的下一个值,因此 out_val 被赋值为 None,所以程序执行 yield i**2(此时 i 己经递增到 3,因此生成器返回 9),程序再次被“冻结”。因此看到本次获取生成器的值为 9。

运行上面程序,可以看到如下输出结果:

0
1
--------------
====9
81
9

此外,生成器还提供了如下两个常用方法:
  1. close():该方法用于停止生成器。
  2. throw():该方法用于在生成器内部(yield 语句内)引发一个异常。

例如,在程序中增加如下代码:
# 让生成器引发异常
sg.throw(ValueError)
运行上面代码,将看到如下输出结果:

Traceback (most recent call last):
  File "C:\Users\mengma\Desktop\1.py", line 21, in <module>
    sg.throw(ValueError)
  File "C:\Users\mengma\Desktop\1.py", line 6, in square_gen
    out_val = (yield out_val ** 2) if out_val is not None else (yield i ** 2)
ValueError

从上面的输出结果可以看到,在程序调用生成器的 throw() 方法引发异常之后,程序就会在 yield 语句中引发该异常。

将上面的 sg.throw(ValueError) 代码注释掉,为程序增加如下两行代码来示范 stop() 方法的用法。在程序调用 stop() 方法关闭生成器之后,程序就不能再去获取生成器的下一个值,否则就会引发异常。
# 关闭生成器
sg.close()
print(next(sg)) # StopIteration
运行上面代码,可以看到如下输出结果:

Traceback (most recent call last):
  File "C:\Users\mengma\Desktop\1.py", line 24, in <module>
    print(next(sg)) # StopIteration
StopIteration