CODEWITHSUNDEEP
pythonpython-3.x
Armandas
Armandas

Reputation: 2396

How to return a value, but keep executing

I have a simple tree class and I need to be able to iterate over the leaf nodes only.

I add the data to my tree as follows:

p = ParamTree()
p.add('system.screen.timeout')
p.add('system.screen.lock.enabled')
p.add('system.screen.lock.code')

and I want to be able to get timeout, enabled and code in sequence.

If I write a function to print the values, it works as expected:

def print_nodes(tree):
    for node in tree.nodes:
        if node.has_children():
            print_nodes(node)
        else:
            print(node)

Output:

>>> print_nodes(p)
timeout
enabled
code

How would I go about implementing a generator that does the same thing? I tried replacing print() with yield, but it doesn't work.

def yield_nodes(tree):
    for node in tree.nodes:
        if node.has_children():
            yield_nodes(node)
        else:
            yield node

Output:

>>> g = yield_nodes(p)
>>> for n in g:
...   print(n)
...
>>>

Upvotes: 1

Views: 68

Answers (1)

Kendas
Kendas

Reputation: 2243

I'm going to assume your yield_nodes generator to be written as follows:

def yield_nodes(tree):
    for node in tree.nodes:
        if node.has_children():
            yield_nodes(node)
        else:
            yield node

As you might notice, calling yield_nodes returns a generator (that you could iterate over), but you are not actually doing anything with it. My proposed solution would be the following:

def yield_nodes(tree):
    for node in tree.nodes:
        if node.has_children():
            yield from yield_nodes(node)
        else:
            yield node

Upvotes: 2

Related Questions