max max - 7 months ago 25
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

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

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)