使用生成器
生成器与生成器函数
如果一个函数包含 yield
表达式,那么它是一个生成器函数,调用它会返回一个生成器。
生成器也是一种迭代器,在每次迭代的时候返回一个值,直到抛出
StopIteration
异常。
def func():
return 1
def gen():
yield 1
if __name__=='__main__':
print(type(func)) # <class 'function'>
print(type(gen)) # <class 'function'>
print(type(func())) # <class 'int'>
print(type(gen())) # <class 'generator'>
可以看到生成器函数与普通函数的区别在于:
- 生成器函数没有
return
,而是yield
- 生成器函数返回的是一个生成器对象
那么它与一般的函数有何不同呢,它的特点在哪呢?下面就简单讲两个例子。
例一、读取文件
通常我们可以使用生成器来读取文件,比如csv
文件。现在,我们想要计算csv文件的行数时,可以分别使用如下两种代码来完成,结果如下:
from memory_profiler import profile
@profile
def csv_func(file_path):
file = open(file_path)
csv_gen = file.read().split("\n")
row_count = 0
for row in csv_gen:
row_count += 1
print(f"Row count is {row_count}")
csv_func('E:\\pycharmProject\\webserverProject\\techcrunch.csv')
这份代码中,我们打开文件,并将数据读取到列表csv_gen
中,然后每读取一个元素,row_count
就加一。然后,我们使用生成器来达到我们的目标,代码如下:
def csv_test(file_path):
for row in open(file_path, "r"):
yield row
@profile
def csv_generator(file_path):
csv_gen = csv_test(file_path)
row_count = 0
for row in csv_gen:
row_count += 1
print(f"Row count is {row_count}")
csv_generator('E:\\pycharmProject\\webserverProject\\techcrunch.csv')
这两份代码中,我们还导入了memory_profiler
库,来查看运行中的内存的使用情况。那么结果如下:
#以下是用列表的方式读取数据的结果
Row count is 504051
Line # Mem usage Increment Line Contents
================================================
9 16.5 MiB 16.5 MiB @profile
10 def csv_func(file_path):
11 16.5 MiB 0.0 MiB file = open(file_path)
12 79.7 MiB 63.2 MiB csv_gen = file.read().split("\n")
13 79.7 MiB 0.0 MiB row_count = 0
14
15 79.7 MiB 0.0 MiB for row in csv_gen:
16 79.7 MiB 0.0 MiB row_count += 1
17
18 79.7 MiB 0.0 MiB print(f"Row count is {row_count}")
#以下是用生成器读取数据的结果
Row count is 504050
Line # Mem usage Increment Line Contents
================================================
20 16.5 MiB 16.5 MiB @profile
21 def csv_generator(file_path):
22 16.5 MiB 0.0 MiB csv_gen = csv_test(file_path)
23 16.5 MiB 0.0 MiB row_count = 0
24
25 16.7 MiB 0.1 MiB for row in csv_gen:
26 16.7 MiB 0.0 MiB row_count += 1
27
28 16.7 MiB 0.0 MiB print(f"Row count is {row_count}")
对比结果,看第一部分的12行csv_gen = file.read().split("\n")
,说明用列表读取数据的话,内存占用增量为63.2MiB;看第二部分25行,可以知道用生成器读取数据的话,内存的占用增量为0.1MiB,显然,使用生成器可以有效减少内存占用。
其实也很好理解,当我们存在一个容器时,想要遍历其中的值,有两种做法:
- 先将容器中的所有值都取出来,然后进行遍历
- 从头开始,边取值边遍历
显然,第二种是更加节省内存空间的。
实际上,当文件大到一定程度时,只有第二种代码才能读取到数据内容,使用第一种代码的话,会报错MemoryError
。
例二、产生一个无限序列
首先,先着眼于有限序列,我们可以使用range()
函数来生成一个有限序列。
>>> a = range(5)
>>> list(a)
[0, 1, 2, 3, 4]
然而,要产生一个无限序列的话,我们就需要使用到生成器了,因为我们的内存是有限的。代码如下:
import time
def infinite_sequence():
num=0
while True:
yield num
num+=1
if __name__=='__main__':
for i in infinite_sequence():
print(i,end=' ')
time.sleep(0.1)
代码逻辑很简单:当infinite_sequence()
函数执行时,yield
生成一个数值并返回,且保留函数当时的num值,然后生成下一个数。所以通过它可以产生一个无限序列。
也可以不使用for
循环,而是先通过生成器函数得到生成器,然后调用next()
方法作用于生成器对象,来获取下一个生成的值。
>>> def infinite_sequence():
... num=0
... while True:
... yield num
... num+=1
...
>>> gen=infinite_sequence()
>>> type(gen)
<class 'generator'>
>>> next(gen)
0
>>> next(gen)
1
>>> gen.__next__()
3
>>> gen.__next__()
4
那么简单分析一下:
- 调用
infinite_sequence()
函数不会立即执行代码,而返回了一个生成器对象gen
- 当使用
next()
(在for
循环中会自动调用next()
) 作用于返回的生成器对象时,函数开始执行,在遇到 yield 的时候会【暂停】,并返回当前的迭代值; - 当再次使用
next()
的时候,函数会被唤醒,从原来【暂停】的地方继续执行,直到遇到yield
语句,如果没有yield
语句,则抛出异常; - 当使用 yield 时,它会自动创建
__iter__()
和__next()__
方法,即简单高效地生成了迭代器,gen.__ next__()
返回值和使用next()的效果相同也说明了这一点。
生成器函数【暂停】时,会保留当时的上下文环境(即位置和变量);
小结
结合两个例子,总结一下:
- yield语句把函数变为一个生成器函数,函数返回值是生成器。
- 生成器函数通过yield语句可以简洁地生成
__iter__()
和__next()__
方法,即简洁地生成一个迭代器。 - 相比于一般函数,使用生成器可以节省内存开销 。
- 生成器函数的执行过程看起来就是不断地
执行->中断->执行->中断
的过程。yield
语句暂停函数返回迭代值,next()
语句唤醒函数继续执行。