2020-08-21

Recursing with yield

This post shows a simple python/js pattern for recursing through code that probably should get a bit more usage.

Say you have the following data:

tree = {
    "a": 4,
    "children": [
        {
            "a": 3,
            "children": []
        },
        {
            "a": 2,
            "children": [
                {
                    "a": 6,
                    "children": []
                }
            ]
        },
    ]
}

And you want to get all the a values greater than 2.

In python, the "classic" way would be to do:

def greater_than_2(t):
    a, children =  t["a"], t["children"]
    accumulated = ()
    if a > 2:
        accumulated += (a, )
    for child in children:
        accumulated += greater_than_2(child)
    return accumulated

greater_than_2(tree)

That will give you: (4, 3, 6).

I find it is often more natural to use yield:

def greater_than_2(t):
    a, children =  t["a"], t["children"]
    if a > 2:
        yield a
    for child in children:
        yield from greater_than_2(child)

tuple(greater_than_2(tree))

A near identical solution can be achieved with modern js (this may surprise any python devs that haven't seen the language for a while 😊 ):

function* greater_than_2(t){
    const {a, children} = t
    if (a > 2) yield a
    for (const child of children) yield* greater_than_2(child)
}

[...greater_than_2(tree)]