Using generator send() within a for loop

Question

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 send() method), and if it's True, 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)
try:
  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:
    pass

Is there any way to improve this?

I thought something like this should work:

for n in bfs(g, start_node):
  ???.send(process(n))

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

Answer 1

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):
    try:
      x = next(generator)
      while True:
        yield x
        x = generator.send(follow_function(x))
    except StopIteration:
        pass

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