How to stop a recursive generator in Python?-CodePudding

Let's say I'm using a recursive generator to traverse a binary tree and I'd like the traversal to stop when a certain condition is satisfied:

def traverse(node):
    if node is None:
        return

    yield from traverse(node.left)

    if node.val == 5:
        return # return only skips node 5, but doesn't fully stop iteration. What does?

    yield node.val

    yield from traverse(node.right)

How do you achieve this?

I tried raising StopIteration, but that results in a RuntimeException.

CodePudding user response：

Separate stop and iteration into two:

from itertools import takewhile

def traverse(node, stop_value):
    yield from takewhile(lambda x: x != stop_value, _traverse(node))

def _traverse(node):
    if node is None:
        return
    yield from traverse(node.left)
    yield node.val
    yield from traverse(node.left)

CodePudding user response：

You should control when the iteration stops from whatever code calls traverse():

def traverse(node):
    if node is None:
        return

    yield from traverse(node.left)

    yield node.val

    yield from traverse(node.right)

for n in traverse(node):
    if n == 5:
        break