Dictionary leaf generator

December 13, 2021

Here it is my tree, a nested dictionary

tree = {"root":
           {"branch_a": {"branch_aa": 0, "branch_ab": 1},
            "branch_b": 1,
            "branch_c": {"branch_ca": {"branch_caa": 0}}}}

I managed to write a function that prints all the leaves

def print_leaves(tree):
    if not hasattr(tree, "__iter__"):
        print(tree)
    elif isinstance(tree, dict):
        for branch in tree.values():
            print_leaves(branch)

which produces the desired output

At this point I thought it would have been nice to decouple the action (printing in this case) from the access to the leaves. So I modified the function above slightly, turned it into a generator and moved the printing part in a for loop.

def generate_leaves(tree):
    if not hasattr(tree, "__iter__"):
        yield tree
    elif isinstance(tree, dict):
        for branch in tree.values():
            generate_leaves(branch)

for leaf in generate_leaves(tree):
    print(leaf)

… which unfortunately doesn’t work.

First of all why, doesn’t it work?
And then, of course, how to properly write a leaf generator?

>Solution :

You are not using the result of the recursive call. That worked for print_leaves which didn’t have a return value, but that doesn’t work for a function with return or yield.

Here is the long version:

def generate_leaves(tree):
    if not hasattr(tree, "__iter__"):
        yield tree
    elif isinstance(tree, dict):
        for branch in tree.values():
            for leaf in generate_leaves(branch):
                yield leaf

for leaf in generate_leaves(tree):
    print(leaf)

Fortunately, we can make that a lot shorter with yield from:

def generate_leaves(tree):
    if not hasattr(tree, "__iter__"):
        yield tree
    elif isinstance(tree, dict):
        for branch in tree.values():
            yield from generate_leaves(branch)

for leaf in generate_leaves(tree):
    print(leaf)

Note that you only need an if / else with one condition; no need for an if / elif with two redundant conditions:

def generate_leaves(tree):
    if not isinstance(tree, dict):
        yield tree
    else:
        for branch in tree.values():
            yield from generate_leaves(branch)

for leaf in generate_leaves(tree):
    print(leaf)

Resources about yield from: