【Python测试开发系列】尽量使用生成器

通过列表生成式，我们可以直接创建一个list，如果生成list长度非常大，是非常消耗内存的。比如，0-9999的平方组成的列表[x**2 for x in range(10000)]占用的内存空间是87624字节（通过sys.getsizeof()获得）。如果创建一个100万或者1000万个元素的list，占用的内存空间那将会更大。

在循环使用元素的过程中不断推算出后续的元素，这样就不必在使用元素之前，就创建完整的list，就能节省大量的内存空间。

在Python中，一边使用一边推算后面元素值的机制叫生成器Generator。

生成器是一种特殊的迭代器。他的特殊语法能让我们不用定义类去实现__iter__() and __next__()就能得到一个迭代器。这个特殊的语法就是yield。

1 如何定义一个生成器

1.1 最简单的生成器

Python中定义生成器有两种方式，生成器表达式 and 生成器函数。生成器表达式的例子如下：

numbers = [1, 2, 3, 4, 5, 6]
lazy_squares = (x * x for x in numbers)

可见，generator占用的空间是非常小的。想要获取lazy_squares中的元素，可以通过next()

print(next(lazy_squares))  # 输出: 1
print(next(lazy_squares))  # 输出: 4
print(next(lazy_squares))  # 输出: 9
print(next(lazy_squares))  # 输出: 16
print(next(lazy_squares))  # 输出: 25
print(next(lazy_squares))  # 输出: 36

generator保存的是算法，每次调用next()，就计算出下一个元素的值，直到计算到最后一个元素，没有更多的元素时，抛出StopIteration的异常。我们可以使用for循环来读取生成器中的元素：

for i in g:
    print(i)

注意，生成器只能被迭代一次。只能往前迭代不能往后。

1.2 使用yield定义生成器

一个带有 yield 语句的函数就是一个 生成器函数。yield语句的作用是返回某个值，并暂停程序的执行。下面的代码定义了一个名为counter的生成器。

def counter():
    print("phase-1")
    yield 1
    print("phase-2")
    yield 2
    print("phase-3")
    yield 3
g=counter() # 生成一个 generator对象 ,不会执行counter中的任何函数代码
print(next(g)) # 函数中的代码开始执行，输出phase-1和1，并且停在第一个yield处
print(next(g))  # 从上一个yield的下一条语句开始执行，输出phase-2和2，遇到第二个yield时代码挂起，返回到next函数调用处
print(next(g))  # 从上一个yield的下一条语句开始执行，输出phase-3和3，停在第三个yield处
print(next(g))  # 没有更多的元素时，抛出StopIteration的异常

代码的执行过程，在代码的注释中已经写得很清楚了。可以看到，在对generator调用 next()方法时，函数中的代码才开始执行。每执行到一个 yield 语句就会中断，并返回一个迭代值给next函数，下次执行时从 yield 的下一个语句继续执行。

第一个next调用相当于启动生成器，会从生成器函数的第一行代码开始执行，直到第一次执行完yield语句，之后跳出生成器函数。

然后第二个next调用，进入生成器函数后，从yield语句的下一句语句开始执行，运行到yield语句，之后跳出生成器函数，后面再次调用next，依次类推。

使用yield定义生成器，是最通用的做法，下面再继续举几个例子。

例子一：倒计时数器

generator适合返回那些使用某种算法推算的情况，下面是一个简单的推算算法，下一个元素比前一个元素少1。

def countdown(n):
    while n > 0:
        yield n  # 返回一个n值
        n -= 1  # 根据推算算法推算下一个元素
c=countdown(3)
print("next函数得到返回值=",c.__next__())
print("next函数得到返回值=",c.__next__())
print("next函数得到返回值=",c.__next__())

例子二：斐波拉契数列

这个要比前面的那个推算规则复杂一些。后一个元素是前两相邻两个元素的和。

def fib(bound):
    n = 0
    a, b = 0, 1
    while n < bound:
        yield a  # 生成一个值
        a, b = b, a + b  # 根据推算算法推算下一个元素
        n += 1

for i in fib(19):
    print(i)

通过前面两个例子，可以看出使用yield定义generator的套路，就是在一个循环中，yield一个值，然后在下一条语句，写推算算法。

3 生成器的高级方法

除了__next__方法，用户还可以通过生成器的send方法启动生成器，并给生成器的yield表达式赋值。其实__next__()和send()在一定意义上作用是相似的，区别是send()可以传递yield表达式的值进去，而__next__()不能传递特定的值，只能传递None进去。

因此，我们可以看做c.__next__() 和 c.send(None) 作用是一样的。需要提醒的是，第一次调用时，请使用__next__()语句或是send(None)，不能使用send发送一个非None的值，否则会出错的，因为没有Python yield语句来接收这个值。

send函数主要是用于外部与生成器对象的交互。下面看一个例子

def gen(x):
    count = x
    while True:
        val = (yield count) # (yield count)会接收到send发送过来的值
        if val is not None:
            count = val
        else:
            count += 1
f = gen(5)
print(f.__next__()) 
print(f.__next__()) 
print(f.send(100)) #在外面控制生成器的执行
print(f.__next__()) 
print(f.__next__()) 

上面代码的输出结果是：

5
6
100
101
102

上面代码的执行流程如下：

1. 调用gen()函数产生一个生成器，此时没有执行gen()函数里的任何代码
2. 第一次调用生成器的__next__方法会启动生成器，进入gen()里，从第一行代码执行到(yield count)处停止，将5返回给__next__函数。此时变量val还没有创建，所以val=None。
3. 第二次调用__next__，先将语句块(yield count)的返回值赋给val，即None给val。然后执行语句块(if val is not None)，count变成6，进入下一次循环，再次执行到yield，返回6。
4. 调用send方法将生成器中 (yield count) 表达式的值赋为send方法的参数，即 (yield count) =100，代码继续执行，count被赋值为100，直到遇到下一个yield语句将100返回给send函数。
5. 再次调用__next__，(yield count)将被赋值为None，导致val=None，代码继续执行，count加1变成101，直到遇到下一个yield语句将101返回给__next__函数

这个挺难理解的，多执行几次看看。

4 生成器的实际用处

生成器能够让你写出简洁的代码，使用更少的中间变量，并且非常节省CPU和内存。
计算一个列表中每个元素的平方的代码如下：

def something(param):
    result = []
    for x in param:
        result.append(x ** 2)
    return result

if __name__ == '__main__':
    print(something([1, 2, 3]))

使用生成器，改写后的代码是：

def iter_something(param):
    for x in param:
        yield x**2

if __name__ == '__main__':
    for i in iter_something([1,2,3]):
        print(i)

下面的代码相比较上面的代码，有以下优点：

• 不需要中间变量result，将中间放进去，最后再返回。
• 节省了空间，不管param参数多大，生成器都不需要缓存整个内容。
• 可以处理无限流，比如stdin、kafka的consumer、文件流、socket等。
• 更快。不需要整个输入参数全部处理完成后再返回，而是一边处理一边输出。
• 由使用者决定如何使用结果。

最后一个优点，是非常重要的。使用者可以根据需要对结果进行去重、取最大值、或者过滤出部分数据。上面的例子是进行了迭代打印。

uniq = set(iter_something([1, 2, 3, 4, 5, 3]))  # 去重
lagest = max(iter_something([1, 2, 3, 4, 5, 3]))  # 取最大值
slicei = list(islice(iter_something([1, 2, 3, 4, 5, 3]), 1, 4))  # 切片

4.1 生产者与消费者并发

利用send方法来控制generator的执行，可以实现在单线程的情况下实现并发运算，实现了生产者和消费者的并发，明显地看到两个任务的打印是你一次我一次，即并发执行的。

def consumer():
    r = ''
    while True:
        n = yield r
        print("[Consumer] 收到第{}个包子 @{}".format(n, int(time.time())))
        if not n:
            return
        print("[Consumer] 吃包子 {}...".format(n))
        r = '我吃完了'

def produce(c): 
    c.send(None)  # 1、启动generator，consumer从代码第一行执行到yield停下来，接着执行produce()中接下来的代码
    x = 0
    while x < 3:
        x += 1
        print("[Producer] 生产第 %s 个包子, @%s" % (x, int(time.time())))
        s = c.send(x)  # 2、给consumer的yield表达式传递参数x，consumer接着上一次暂停处往下继续执行，r最终变成"我吃完了"，接着循环遇到yield，consumer()函数又暂停并且返回变量r的值，此时程序又进入produce(c)函数中接着执行，s收到consumer yield出来的值
        print("[Producer] 顾客说: %s" % s)
    c.close()

cc = consumer() # 生成一个 generator ,这里不会执行任何函数代码
produce(cc) # 进入produce函数执行c.send(Node)

代码输出：

[Producer] 生产第 1 个包子, @1563861836
[Consumer] 收到第1个包子 @1563861836
[Consumer] 吃包子 1...
[Producer] 顾客说: 我吃完了
[Producer] 生产第 2 个包子, @1563861836
[Consumer] 收到第2个包子 @1563861836
[Consumer] 吃包子 2...
[Producer] 顾客说: 我吃完了
[Producer] 生产第 3 个包子, @1563861836
[Consumer] 收到第3个包子 @1563861836
[Consumer] 吃包子 3...
[Producer] 顾客说: 我吃完了

上面代码通过send控制generator的执行过程，具体执行流程如下：

1. 调用consumer()函数产生一个generator，把一个generator传入produce；
2. produce函数中，首先调用c.send(None)启动生成器，consumer函数开始执行，会在yield处停止，接着返回到produce函数中，从x=0这一句往下执行；
3. produce函数执行到c.send(x)，将x的值传递到consumer中的yield表达式（yield r），n就等于x了，接着执行第一条print，执行判断，执行第二条print，给r赋值，又回到yield语句，把r返回给send函数，从而传递给s；
4. produce拿到consumer处理的结果，继续生产下一条消息；
5. x=3时，通过c.close()关闭consumer，整个过程结束。

不过在Python3.6之后，并发编程有更好的选择，就是使用asyncio包。

4.2 提取多个城市的天气信息

定义一个generator，接收一个城市列表，在for循环中每次查询一个城市的数据。这样避免在get_weather_7_days函数中，使用了列表保存城市的天气数据了。

import requests

def get_weather_7_days(citys):
    url = "https://www.tianqiapi.com/api/?version=v1&city={city}"
    headers = {
        "content-type": "application/x-www-form-urlencoded"
    }
    for city in citys:
        yield requests.get(url.format(city=city), headers=headers).json()

for weather in get_weather_7_days(["北京", "天津", "青岛"]):
    print(weather)

与这个例子类似，使用yield定义爬虫的生成器也是常用的方法。

4.3 监控日志中的500错误。

因为日志文件很难知道什么时候结束，我们读取文件内容时，并不能一次性读到所有的内容。对于这样的情况非常适合使用generator。我们先写一个监控脚本monitor.py如下:

import time


def tail(file_path):  # 定义一个查看文件的函数
    with open(file_path, 'r') as f:
        f.seek(0, 2)  # 把游标移动到文件末尾
        while True:  # 循环监控日志
            data = f.readline()  # 读取文件末尾
            if data:  # 如果有数据就用yield返回
                yield data
            else:
                time.sleep(1)

def grep(file, k):  # 定义过滤关键字函数
    for i in tail(file):  # 循环生成器中的数据
        if k in i:
            print(i)

if __name__ == '__main__':
    grep('a.txt', '500')  # 监控a.txt最新日志，并过滤500的错误代码，一旦有500出现就会被抓到

我们再模拟一个写日志的脚本writer.py：

import time
import random

with open('a.txt', 'a') as f:
    while True:
        f.write("500\n")
        f.flush()
        time.sleep(random.randint(1, 10))

将这两个脚本分别启动起来，可以看到当writer.py写入500后，monitor.py会输出出来。

5.总结

生成器是一种特殊的迭代器，可以用for循环来进行迭代。生成器能够节省内存，在迭代的过程中不断输出新的元素。可以实现一种懒加载的机制，在使用元素时才生成元素。
生成器的实际应用还是比较广泛的，目的主要是懒加载和节省内存。如果你想对一个函数的输出进行迭代，那么第一想到的是将这个函数实现成一个迭代器。

参考文章

1. http://www.cnblogs.com/xybaby/p/6322376.html
2. https://www.liaoxuefeng.com/wiki/0014316089557264a6b348958f449949df42a6d3a2e542c000/001432090954004980bd351f2cd4cc18c9e6c06d855c498000
3. https://www.jb51.net/article/113160.htm
4. https://segmentfault.com/a/1190000016880292
5. https://www.bbsmax.com/A/ZOJPQjq25v/
6. https://nvie.com/posts/iterators-vs-generators/
7. https://nvie.com/posts/use-more-iterators/