max max - 1 year ago 59
Python Question

Using generator send() within a for loop

I implemented graph traversal as a generator function which yields the node being visited.

Sometimes the user needs to tell the traversal function that the edges outgoing from a particular node shouldn't be followed; in order to support that, the traversal checks the value sent back to it (using generator

method), and if it's
, regards the node as a leaf for traversal purposes.

The problem is that the simplest user loop is kinda long:

# simplified thanks to @tobias_k
# bfs is the traversal generator function
traversal = bfs(g, start_node)
n = next(traversal)
while True:
# process(n) returns True if don't want to follow edges out of n
n = traversal.send(process(n))
except StopIteration:

Is there any way to improve this?

I thought something like this should work:

for n in bfs(g, start_node):

but I feel I'm missing the knowledge of some python syntax.


I don't see a way to do this in a regular for loop. However, you could create another generator, that iterates another generator, using some "follow-function" to determine whether to follow the current element, thus encapsulating the tricky parts of your code into a separate function.

def checking_generator(generator, follow_function):
      x = next(generator)
      while True:
        yield x
        x = generator.send(follow_function(x))
    except StopIteration:

for n in checking_generator(bfs(g, start_node), process):