Python generator和yield 的使用

Python generator和yield 的使用

Jan 17, 2014
Coding
Python

本文内容适用于 python 2.x 版本

Python中的generator和yield关键字用来实现类似iterator迭代器的功能。

使用了yield关键字的函数会被视为一个generator，一个generator就类似于iterator，是可以迭代的对象。

generator执行到yield语句时会返回yield语句返回的值，然后保存的当前状态（包括函数中各种变量的当前值），等到下一次调用generator的next()时，从yield后面的语句继续执行。

一个简单示例 # 

def foo():
    yield 1
    yield 2
    yield 3

a = foo()
print a.next() #1
print a.next() #2
print a.next() #3

foo()函数中使用yield关键字，表明这个函数是一个generator，其中使用了三次yield分别返回1、2、3。将foo()赋值给a，则a成为了一个generator对象。一个generator对象的功能和iterator对象类似。可以调用a.next()来获取下一个迭代值：

第一次执行a.next()时，foo()运行到yield 1，返回1，然后保存foo()的当前状态，等待下一次调用next()。此时print语句打印出1

第二次执行a.next()时，foo()恢复上次保存的状态，从上次返回的yield语句后继续执行，即执行yield 2。返回2，然后再次保存当前状态，等待下一次调用next()。此时print语句打印出2

第三次执行a.next()时，同理，打印出3。

如果再调用一次a.next()则会出错，产生一个StopIteration异常，因为foo()函数已经执行return返回了。

for循环可以隐式调用generator对象的next() # 

def foo():
    yield 1
    yield 2
    yield 3

for i in foo():
print i #1 2 3

for 循环中，每次隐式的调用generator对象foo()的next()，将返回值赋值给i。for循环会自动处理next的调用，不会产生StopIteration异常。

一个稍复杂的示例 #

generator和yield可以实现迭代器，Python中实现了__iter__()方法和next()方法的对象都是可迭代的迭代器对象。那么两者有什么区别？

使用generator和yield更加简洁

比如要实现打印N以内的水仙花数的程序。

水仙花数是指一个 n 位数 ( n≥3 )，它的每个位上的数字的 n 次幂之和等于它本身。（例如：1^3 + 5^3 + 3^3 = 153）

使用generator和yield实现的水仙花数代码

#coding=utf-8
#author=JarvisChu

def isWaterFlower(num):
    '''判断n是不是水仙花数'''
    s = str(num)
    n = len(s) #n位
    temp_sum = 0
    for i in s:
        temp_sum += int(i)**n #
    return temp_sum == num

def waterflower(N):
    '''返回100 – N的水仙花数'''
    start = 100
    while start <= N:
        if isWaterFlower(start):
            yield start #返回一个水仙花数
        start += 1

for i in waterflower(1000):
    print i #153 370 371 407

waterflower是一个generator对象，每次迭代返回100-N的水仙花数中的下一个。代码简洁明了。

使用__iter__和next()实现的代码

#coding=utf-8
#author=JarvisChu
class WaterFlowerNumber(object):
    def __init__(self,N):
        self.UpBound = N
        self.start = 100 #水仙花数至少是3位

    def __iter__(self):
        return self

    def isWaterFlower(self,num):
        '''判断n是不是水仙花数'''
        s = str(num)
        n = len(s) #n位
        temp_sum = 0
        for i in s:
            temp_sum += int(i)**n #
        return temp_sum == num

    def next(self):
        while self.start <= self.UpBound:
            if self.isWaterFlower(self.start):
                self.start += 1
                return self.start-1 #返回一个水仙花数
            else:
                self.start += 1
            else:
                raise StopIteration()

for i in WaterFlowerNumber(1000):
    print i #153 370 371 407

需要实现__iter__和next()，相比较而言，代码就繁琐的多。

参考阅读 #