Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python generators: correct code recursing a tree

class Node(object):
    def __init__(self, lst):
        if type(lst) == list:
            self.value = lst[0]
            self.children = lst[1:]
        else:
            self.value = lst
            self.children = []
    @property
    def ChildElements(self):
        return [Node(a) for a in self.children]

    @property
    def GetValue(self):
        return self.value

def node_recurse_generator(node):
    yield node.value
    for n in node.ChildElements:
        node_recurse_generator(n)

Node is a simple tree like data structure. The list's first element is always the value of the Node, further elements will be children. If the Node is initiated with something other than a list, that value will be that, and children will be [], an empty list.

a = Node([1,[10,20,30],[100,200,300]])
>>> list(node_recurse_generator(a))
[1]

Now, it would be great to have a recursive iterator over all the elements, but my generator only yields the first element. Why is this happening?

like image 283
George Avatar asked Oct 03 '11 11:10

George


2 Answers

Simply calling node_recurse_generator recursively isn't enough - you have to yield its results:

def node_recurse_generator(node):
    yield node.value
    for n in node.ChildElements:
        for rn in node_recurse_generator(n):
            yield rn
like image 112
NPE Avatar answered Oct 03 '22 22:10

NPE


As yak mentioned in the comment of the top answer, you can also use yield from after Python 3.3.

def node_recurse_generator(node):
    yield node.value
    for n in node.ChildElements:
        yield from node_recurse_generator(n)
like image 30
Zhe He Avatar answered Oct 03 '22 21:10

Zhe He