在Python中从块内停止生成器

Question

我有一个生成器，它可以从一个有向无环图（DAG）中按深度优先的方式输出节点：

def depth_first_search(self):
    yield self, 0 # root
    for child in self.get_child_nodes():
        for node, depth in child.depth_first_search():
            yield node, depth+1

我可以这样遍历这些节点

for node, depth in graph.depth_first_search():
    # do something

我希望能够在循环中告诉生成器，如果满足某个条件，就停止深入图中。

我想出了一个解决方案，使用了一个外部函数。

def depth_first_search(self, stop_crit=lambda n,d: False):
    yield self, 0 # root
    for child in self.get_child_nodes():
        for node, depth in child.depth_first_search():
            yield node, depth+1
            if stop_crit(node, depth): break

这个解决方案让我必须在定义 stop_crit 之前声明我需要的变量，这样才能在里面访问到它们。

在 Ruby 中，yield 会返回块中的最后一个表达式，所以可以方便地用来告诉生成器继续还是停止。

在 Python 中，怎么才能实现这样的功能呢？

Answer 1

协程（bassfriend提到过）对于初学者来说有点复杂，所以这里给大家介绍一个例子。我添加了一些测试代码，这样你可以看到它是如何工作的。

class Node(object):
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

    # the producing coroutine, it sends data to the consumer
    def depth_first_search(self, consumer, depth=0):
        """ `consumer` is a started coroutine that yields True to continue a branch """
        if consumer.send((self, depth)): # continue this branch?
            for child in self.get_child_nodes():
                child.depth_first_search(consumer, depth+1)

    def get_child_nodes(self):
        for node in (self.left, self.right):
            if node is not None:
                yield node

    def __repr__(self):
        return "Node(val=%d)" % self.val

def coroutine(func):
    """ decorator that declares `func` as a coroutine and starts it """
    def starter(*args, **kwargs):
        co = func(*args, **kwargs)
        next(co) # corotines need to be started/advanced to the first yield
        return co
    return starter

# the consumer takes data and yields if it wants to continue
@coroutine
def consumer( continue_branch=lambda n,d:True ):
    node, depth = (yield True) # first node
    while True:
        print node, depth # do stuff
        node, depth = (yield continue_branch(node, depth))


# testing
tree = Node(5, Node(2, Node(3), Node(4)), Node(6, Node(7), Node(8))) # 
cons = consumer()
tree.depth_first_search(cons)# yields all

print
stopper = consumer(lambda n,d: n.val > 2) # skips the children of Node 2
tree.depth_first_search(stopper)

关键在于，如果你保持函数的角色，比如depth_first_search这个函数负责生成节点，你最终会搞得一团糟……相反，节点是被生成出来的，然后发送给使用它的地方。

Python对协程的支持有点别扭（@coroutine来帮忙）。有一个很不错的Python教程，还有很多关于依赖协程的语言（比如Lua）的资源。总之，这是一个非常酷的概念，值得深入了解 :-)

Answer 2

通常在Python中，你只需要停止使用生成器，然后就可以不再关注它了。这样的话，Python会自动处理这些不再使用的东西。

不过，如果你想要立即清理生成器，可以使用 generator.close() 这个方法，这样可以立刻结束生成器的工作，并触发所有的清理操作。

举个例子：

>>> def gen():
...     try: 
...         for i in range(10):
...             yield i
...     finally:
...         print "gen cleanup"
...         
>>> g = gen()
>>> next(g)
0
>>> for x in g:
...     print x
...     if x > 3:
...         g.close()
...         break
...        
1
2
3
4
gen cleanup
>>> g = gen()
>>> h = g
>>> next(g)
0
>>> del g
>>> del h   # last reference to generator code frame gets lost
gen cleanup

Answer 3

简单的解决方案：

def depth_first_search(self):
    yield self, 0 # root
    for child in self.get_child_nodes():
        for node, depth in child.depth_first_search():
            if(yield node, depth+1):
                yield None # for .send
                return

你可以像平常一样调用它，但你需要保存这个可迭代对象，以便可以中止：

it = graph.depth_first_search()
for node, depth in it: #this is why there should be pronouns for loop iterables
    stuff(node,depth)
    if quit: it.send(1) 
    # it.next() should raise StopIteration on the next for iteration

我觉得这个方法现在是有效的。