在python中怎样去使用生成器与yield函数（二）生成器

使用生成器

生成器与生成器函数

如果一个函数包含 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'>

可以看到生成器函数与普通函数的区别在于：

1. 生成器函数没有return，而是yield
2. 生成器函数返回的是一个生成器对象

那么它与一般的函数有何不同呢，它的特点在哪呢？下面就简单讲两个例子。

例一、读取文件

通常我们可以使用生成器来读取文件，比如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，显然，使用生成器可以有效减少内存占用。
其实也很好理解，当我们存在一个容器时，想要遍历其中的值，有两种做法：

1. 先将容器中的所有值都取出来，然后进行遍历
2. 从头开始，边取值边遍历

显然，第二种是更加节省内存空间的。
实际上，当文件大到一定程度时，只有第二种代码才能读取到数据内容，使用第一种代码的话，会报错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

那么简单分析一下：

1. 调用 infinite_sequence()函数不会立即执行代码，而返回了一个生成器对象gen
2. 当使用 next() (在 for 循环中会自动调用 next()) 作用于返回的生成器对象时，函数开始执行，在遇到 yield 的时候会【暂停】，并返回当前的迭代值；
3. 当再次使用 next() 的时候，函数会被唤醒，从原来【暂停】的地方继续执行，直到遇到 yield 语句，如果没有 yield 语句，则抛出异常；
4. 当使用 yield 时，它会自动创建__iter__()和__next()__方法，即简单高效地生成了迭代器，gen.__ next__()返回值和使用next()的效果相同也说明了这一点。

生成器函数【暂停】时，会保留当时的上下文环境（即位置和变量）；

小结

结合两个例子，总结一下：

1. yield语句把函数变为一个生成器函数，函数返回值是生成器。
2. 生成器函数通过yield语句可以简洁地生成__iter__()和__next()__方法，即简洁地生成一个迭代器。
3. 相比于一般函数，使用生成器可以节省内存开销 。
4. 生成器函数的执行过程看起来就是不断地 执行->中断->执行->中断 的过程。yield语句暂停函数返回迭代值，next()语句唤醒函数继续执行。