Stop generator from within block in Python

Question

I have a generator that yields nodes from a Directed Acyclic Graph (DAG), depth first:

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

I can iterate over the nodes like this

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

I would like to be able to tell the generator, from the for loop, to stop from going deeper in the graph if some condition is met.

I came up with the following solution, that use an external function.

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

This solution forces me to declare variables I need before stop_crit is defined so they can be accessed from it.

In Ruby, yield returns the last expression from the block so this could conveniently be used to tell the generator to continue or stop.

What is the best way to achieve this functionality in Python?

Answer 1

Usually in Python you would just stop consuming the generator and forget about it. Point. (Thus leaving things to the garbage collector the usual way)

Yet by using generator.close() you can force an immediate generator cleanup triggering all finalizations immediately.

Example:

>>> 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 2

Naive solution:

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

You can call it normally still, but you have to save the iterable to abort:

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

I think this works right now.