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
0
1
1
0
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?
CodePudding user response:
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
: