理解python生成器

先看一个函数:
[ccb_python]
def group(seq, size):
“””
Returns an iterator over a series of lists of length size from iterable.

>>> list(group([1,2,3,4], 2))
[[1, 2], [3, 4]]
>>> list(group([1,2,3,4,5], 2))
[[1, 2], [3, 4], [5]]
“””
def take(seq, n):
for i in xrange(n):
yield seq.next()

if not hasattr(seq, ‘next’):
seq = iter(seq)
while True:
x = list(take(seq, size))
if x:
yield x
else:
break
[/ccb_python]
函数中没有常见的return,而是使用了yield。yield和return相同的地方是都用来向函数的调用者返回值。下面这个简单的例子展示了yield与return的不同。

[ccb_python]
def my_gen():
“””
A simple generator
“””
print “Generator”
yield 1
yield “The second”
yield 3

if __name__ == ‘__main__’:
g = my_gen()
print “seperator”
print g.next()
print “seperator”
print g.next()
print “seperator”
print g.next()
[/ccb_python]
运行结果:

seperator
Generator
1
seperator
The second
seperator
3

很明显函数中的三次yield都返回了正确的值,并且每次返回之后,函数都会暂停。直到再次调用g.next(),函数又会从上次暂停的地方继续。如果使用return,那么一次return之后控制权就永远回到函数调用者手中。

像上面这样使用yield的函数,都被称作“生成器”。很容易理解,每取走一个值,函数又会继续生成一个新的值。这样的函数就像是不停生产的机器。

与一般的函数不同的是,在调用生成器函数时,其函数体并没有被执行。只有执行next()的时候,函数体才真正开始执行。比如上例中12行调用my_gen函数。但第5行中的print语句并没有执行。而是在第14行过后执行的。下面是一个生成fibonacci数列的生成器。
[ccb_python]
def fib():
a, b = 0, 1
while 1:
yield b
a, b = b, a+b

c = fib()
[/ccb_python]
把生成器函数作为list()、set()、dict()等工厂函数的参数,是很常见的做法。工厂函数本身接受可迭代的对象(iterable)作为参数,并将其转换为对应的类型。生成器本身是可迭代的对象,当然也就可以用于这些工厂函数。比如将某个文件每行第一个字段储存为一个集合,可以这样:
[ccb_python]
def id_gen(path):
for line in open(path):
tmp = line.split()
yield int(tmp[0])

s = set(id_gen(‘path_to_file’))
[/ccb_python]
更pythonic的方式应该是这样:
[ccb_python]
s = set(int(i.split()[0]) for i in open(‘path_to_file’))
[/ccb_python]
同样的效果,只是这个使用了生成器表达式,更加简单明了。python的表达能力很强,有木有?

开头处的take生成器函数接受一个可迭代对象seq和一个整数参数n,返回seq的前n个元素。这样[ccib_python]list(take(seq, size))[/ccib_python]的结果就是由seq的前n个元素所组成的列表(list工厂函数将生成器参数转换为列表)。比如[ccib_python]list(take(iter([1,2,3,4,5]), 2))[/ccib_python]的结果就是[ccib_python][1,2][/ccib_python]这个列表。group生成器先将原始列表参数seq转换为迭代器,然后生成[ccib_python]list(take(seq, size))[/ccib_python]。迭代器seq在整个过程中是一直前进的,所以每次的[ccib_python]list(take(seq, size))[/ccib_python]的结果是原始列表中的下size个元素,而不会始终返回原始列表中的前size个元素。

生成器的应用

那么,生成器有什么用途?仅仅用来一个个算fibonacci数列么?当然不是,生成器的能耐大大超乎你的想象。

生成器给予了python懒惰求值(Lazy evaluation)的能力。通俗的说,就是只有在某个值被用到的时候,它才会被计算出来。比如上面的fib(),这个函数将会产生一个无限的fibonacci数列。那么如果调用它,会把内存消耗地一干二净么?当然不会。因为,只有真正要从里面“取”出一个值的时候,这个值才会被“生产”出来。在当下这个状态,后面那些无穷无尽的值根本还不存在。

从结果来看,也可以认为生成器是返回了多个值,就像返回了一个列表一样。不同的是,列表是在所有值都计算完成后全部返回,而生成器是一个接一个的返回,并且每返回一个值,生成器就暂停,直到它被要求给出下一个值时才继续运行。因为有这个特性,生成器特别适合做一些运算结果大,甚至大到无法分配足够的内存来储存这个结果的任务,比如处理一个超大的文件,把整个文件全部读到内存中是无法实现的,而使用生成器就很容易解决。

一个很好的例子就是os.path.walk()和os.walk()的区别。这两个函数的作用都是遍历给定的目录。前者是在遍历过程中,把所有的数据都放到内存中,直到整个目录都遍历完成之后,才将结果展示出来。后者则使用生成器,每次只返回目录中的一个条目。无论是从内存的使用效率,还是从对用户的体验上讲,后者都有明显的优势。当某个目录中包含相当数量的文件时,第一种方法明显会消耗大量的内存。

在有些情况下,可能需要某个功能函数在执行过程中向它的调用者反馈一些信息。一般的做法是设计一个回调函数,并且将其作为功能函数的参数。需要反馈信息时,功能函数调用回调函数。如果使用生成器,就不再需要回调函数。当它想要返回某些信息时,只要yield就可以了。

生成器还可以把一些非线性的操作“转换”为线性,提高代码的可读性。比如遍历二叉树时,可以先写一个遍历二叉树的生成器函数[ccib_python]bin_gen()[/ccib_python],然后再对节点进行某些操作:[ccib_python]do_something(node) for node in bin_gen()[/ccib_python]。

其实这是利用生成器把控制循环的代码与真正的功能代码分离开来。假如需要处理某些以日期命名的文件,拿过去的n天中任意两天的文件做比较,通常要用两层循环。因为python中的datetime类型不是可迭代类型,由某一天得到下一天的方式只能是:[ccib_python]i = i + datetime.timedelta(1)[/ccib_python]。这样这个两层循环的代码看上去就有点复杂。如果再加上功能部分的代码,那么整个程序看起来就会臃肿不堪。

使用生成器的方案如下:
[ccb_python]
data_dir = “/path/to/data/”

def comp_gen(begin, end):
i = begin
while i + datetime.timedelta(1) < end: x = data_dir + i.strftime('%Y%m%d') j = i + datetime.timedelta(1) while j < end: y = data_dir + j.strftime('%Y%m%d') yield x, y j = j + datetime.timedelta(1) i = i + datetime.timedelta(1) c = comp_gen(begin, end) for i, j in c: do_something(i, j) [/ccb_python] 这样虽然程序的功能相同,但是逻辑变得更清晰易懂。 最后,生成器的“杀手级”应用就是实现协程(coroutine)。Google一下coroutine,出来的结果中第2条与第3条都是python的。想来现在原生支持协程的语言,也就只有python了吧。coroutine其实是同生成器相反的东西,因为它并不是生产数据,而是消耗数据。Coroutine确实是非常非常有用的技术,可惜这里篇幅有限,关于这个神秘的东西,下次再来写吧。

参考:

  1. PEP 255 — Simple Generators Python Generator最基础的介绍
  2. Generator Tricks for Systems Programmers 一个很好的生成器教程,作者提出了一种像shell中的管道一样使用生成器的方法
  3. PEP 342 — Coroutines via Enhanced Generators 用增强的生成器构建协程
  4. PEP 289 — Generator Expressions 生成器表达式
  5. What can you use Python generator functions for? stackoverflow上的一个好问题,生成器可以用来做什么?

发表评论

电子邮件地址不会被公开。 必填项已用*标注