Jamal的博客

Python-生成器

发表于 | 分类于 |

通过List Comprehensions，我们可以很方便的生成一个列表，但是受限于内存等原因，列表的长度是有限制的，例如我们创建一个100w长度的列表，既浪费了空间，又如果我们只需要部分数据的话，那么大部分空间就被浪费了。
所以，如果列表中的元素可以按照某种固定的算法推导出来，我们就可以在迭代中不断地通过这个算法计算出后面的值，这在python中叫做生成器(generator)。

生成器

带有 yield 关键字的的函数在 Python 中被称之为 generator(生成器)。Python 解释器会将带有 yield 关键字的函数视为一个 generator 来处理。一个函数或者子程序都只能 return 一次，但是一个生成器能暂停执行并返回一个中间的结果 —— 这就是 yield 语句的功能 : 返回一个中间值给调用者并暂停执行。
创建一个generator：

1
2
a = [x * x for x in range(10)]
b = (x * x for x in range(10))

可以用过next方法去访问生成器中的元素：

1
2
3
4
5
6
7
8
9
10
11
>>> b = (x * x for x in range(3))
>>> next(b)
0
>>> next(b)
1
>>> next(b)
4
>>> next(b)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

当访问到最后一个元素之后再次访问的时候，会引起StopIteration的报错。
我们来看看next函数：

1
2
3
4
5
6
7
8
def next(iterator, default=None): # real signature unknown; restored from __doc__
    """
    next(iterator[, default])
    
    Return the next item from the iterator. If default is given and the iterator
    is exhausted, it is returned instead of raising StopIteration.
    """
    pass

所以从本质上来说，生成器还是一个迭代器。当然在日常中我们会更多的用for循环来遍历。
前面讲过，实际上生成器保存的是一个算法，而不是一个值。让我们来看下面的一个函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def test():
    print("1")
    yield 1
    print("2")
    yield 2
    print("3")
    yield 3
a = test()
next(a)
print("111")
next(a)
print("222")
next(a)
print("333")

在函数中使用yield，这个函数就变成了一个生成器，我们看一下输出：

1
2
3
4
5
6
1
111
2
222
3
333

会发现实际上是每次都卡住了，我们再看一个斐波那契数列的例子：

1
2
3
4
5
6
7
8
9
10
def fibn(max):
    a, b, n = 0, 1, 0
    while n < max:
        yield b
        a, b = b, a + b
        n += 1
    return 'done'
f = fibn(6)
print(f)

输出：

1
<generator object fibn at 0x1013c00f8>

这里，最难理解的就是generator和函数的执行流程不一样。函数是顺序执行，遇到return语句或者最后一行函数语句就返回。而变成generator的函数，在每次调用next()的时候执行，遇到yield语句返回，再次执行时从上次返回的yield语句处继续执行。
但是在调用生成器的时候，会发现拿不到最后return的值，要获取这个值必须捕获StopIteration的错误，如下代码中：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def fibn(max):
    a, b, n = 0, 1, 0
    while n < max:
        yield b
        a, b = b, a + b
        n += 1
    return 'done'
f = fibn(2)
while True:
    try:
        x = next(f)
        print(x)
    except StopIteration as e:
        print('stopIter error: %s' % e.value)
        break

执行结果：

1
2
3
1
1
stopIter error: done

生成器表达式

生成器表达式是列表解析的拓展，列表解析的不足在于它必须一次性生成所有的数据用于创建对象，所以不适合用于迭代大量的数据。生成器表达式通过结合列表解析和生成器来解决这个问题：

  • 列表解析： [expr for iter_var in iterable if cond_expr]
  • 生成器： (expr for iter_var in iterable if cond_expr)

两者的语法很相似，区别在于生成器返回的是一个生成器对象，而列表解析返回的是一个列表。
我们来看一个读取文件的实例，我们读取一个文件，选择出其中最长的一行并返回长度：

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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
# 代码1: 通用做法,打开文件,通过循环将更长的行赋值给变量longest_line
longest_line = 0
with open('data/data', 'r') as f:
    while True:
        line_len = len(f.readline().strip())
        if not line_len:
            break
        if line_len > longest_line:
            longest_line = line_len
print(longest_line)
# 代码2: 改进1,在读取文件的时候,应该尽早释放文件句柄
longest_line = 0
all_lines = None
with open('data/data', 'r') as f:
    all_lines = f.readlines()
for line in all_lines:
    line_len = len(line.strip())
    if not line_len:
        break
    if line_len > longest_line:
        longest_line = line_len
print(longest_line)
# 代码3: 改进2,使用列表推导来简化代码
longest_line = 0
all_lines = None
with open('data/data', 'r') as f:
    all_lines = [x.strip() for x in f.readlines()]
for line in all_lines:
    line_len = len(line.strip())
    if not line_len:
        break
    if line_len > longest_line:
        longest_line = line_len
print(longest_line)
# 代码4: 改进3,当文件比较大时,直接readlines需要读取文件中所有的行,然后再进行遍历,我们可以使用文件迭代器优化代码
longest_line = 0
all_lines_len = None
with open('data/data', 'r') as f:
    all_lines_len = [len(x.strip()) for x in f]
print(max(all_lines_len))
# 代码5: 改进4,代码4中,在列表推导的时候需要将文件的所有行一下子读取到内存中,然后再创建一个列表对象,我们使用生成器来替代列表推导
longest_line = 0
with open('data/data', 'r') as f:
    longest_line = max(len(x.strip()) for x in f)
print(longest_line)