生成器¶

while 循环通常有这样的形式：

<do setup>
result = []
while True:
    <generate value>
    result.append(value)
    if <done>:
        break

使用迭代器实现这样的循环：

class GenericIterator(object):
    def __init__(self, ...):
        <do setup>
        # 需要额外储存状态
        <store state>
    def next(self): 
        <load state>
        <generate value>
        if <done>:
            raise StopIteration()
        <store state>
        return value

更简单的，可以使用生成器：

def generator(...):
    <do setup>
    while True:
        <generate value>
        # yield 说明这个函数可以返回多个值！
        yield value
        if <done>:
            break

生成器使用 yield 关键字将值输出，而迭代器则通过 next 的 return 将值返回；与迭代器不同的是，生成器会自动记录当前的状态，而迭代器则需要进行额外的操作来记录当前的状态。

对于之前的 collatz 猜想，简单循环的实现如下：

In [1]:

def collatz(n):
    sequence = []
    while n != 1:
        if n % 2 == 0:
            n /= 2
        else:
            n = 3*n + 1
        sequence.append(n)
    return sequence
 
for x in collatz(7):
    print x,

22 11 34 17 52 26 13 40 20 10 5 16 8 4 2 1

迭代器的版本如下：

In [2]:

class Collatz(object):
    def __init__(self, start):
        self.value = start
 
    def __iter__(self):
        return self
 
    def next(self):
        if self.value == 1:
            raise StopIteration()
        elif self.value % 2 == 0:
            self.value = self.value/2
        else:
            self.value = 3*self.value + 1
        return self.value
 
for x in Collatz(7):
    print x,

22 11 34 17 52 26 13 40 20 10 5 16 8 4 2 1

生成器的版本如下：

In [3]:

def collatz(n):
    while n != 1:
        if n % 2 == 0:
            n /= 2
        else:
            n = 3*n + 1
        yield n
 
for x in collatz(7):
    print x,

22 11 34 17 52 26 13 40 20 10 5 16 8 4 2 1

事实上，生成器也是一种迭代器：

In [4]:

x = collatz(7)
print x

<generator object collatz at 0x0000000003B63750>

它支持 next 方法，返回下一个 yield 的值：

In [5]:

print x.next()
print x.next()

22
11

iter 方法返回的是它本身：

In [6]:

print x.__iter__()

<generator object collatz at 0x0000000003B63750>

之前的二叉树迭代器可以改写为更简单的生成器模式来进行中序遍历：

In [7]:

class BinaryTree(object):
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right
 
    def __iter__(self):
        # 将迭代器设为生成器方法
        return self.inorder()
 
    def inorder(self):
        # traverse the left branch
        if self.left is not None:
            for value in self.left:
                yield value
 
        # yield node's value
        yield self.value
 
        # traverse the right branch
        if self.right is not None:
            for value in self.right:
                yield value

非递归的实现：

In [9]:

def inorder(self):
    node = self
    stack = []
    while len(stack) > 0 or node is not None:
        while node is not None:
            stack.append(node)
            node = node.left
        node = stack.pop()
        yield node.value
        node = node.right

In [10]:

tree = BinaryTree(
    left=BinaryTree(
        left=BinaryTree(1),
        value=2,
        right=BinaryTree(
            left=BinaryTree(3),
            value=4,
            right=BinaryTree(5)
        ),
    ),
    value=6,
    right=BinaryTree(
        value=7,
        right=BinaryTree(8)
    )
)
for value in tree:
    print value,

1 2 3 4 5 6 7 8

原文: https://nbviewer.jupyter.org/github/lijin-THU/notes-python/blob/master/05-advanced-python/05.10-generators.ipynb