CODEWITHSUNDEEP
python
robev
robev

Reputation: 1949

Iterate recursion (in a Python Tree)

I've written a Tree class in Python but I'm having issues creating in iterator for it I want to be able to do

phonebook = MyTree()
# Build up tree

for node in phonebook:
    print "%s: %s" % (node.key(), node.data())

But it's not working (says generator object doesn't have key() and data()). My Tree class's __iter__ function returns an iterator class I've created. Here's what I have so far (I know it's wrong and it's not working as it's returning a generator object, since that's what yield does, I want it to remember where it was in the recursion though..so I can't use return) . Basically I just want to return the nodes inorder.

class TreeIterator():
    def __init__(self, root, size):
        self._current = root
        self._size = size
        self.num_visited = 0

    def __iter__(self):
        return self

    def next(self):
        return self._next(self._current)

    def _next(self, curr):
        self.num_visited = self.num_visited + 1
        if self.num_visited == self._size:
            raise StopIteration

        if curr.left is not None and curr.left is not TreeNode.NULL:
            yield self._next(curr.left)

        yield curr

        if curr.right is not None and curr.right is not TreeNode.NULL:
            yield self._next(curr.right)

Upvotes: 2

Views: 3138

Answers (3)

robert
robert

Reputation: 34438

Try changing

if curr.left is not None and curr.left is not TreeNode.NULL:
    yield self._next(curr.left)

yield curr

if curr.right is not None and curr.right is not TreeNode.NULL:
    yield self._next(curr.right)

into

if curr.left is not None and curr.left is not TreeNode.NULL:
    for x in self._next(curr.left):
        yield x

yield curr

if curr.right is not None and curr.right is not TreeNode.NULL:
    for x in self._next(curr.right):
        yield x

It looks like you're yielding an iterator, not a value. I also think your general approach is way too complicated.

self._next(curr.left) returns a generator/iterator. It contains a bunch of values, not just one, so you need to loop through it.

Upvotes: 2

Sven Marnach
Sven Marnach

Reputation: 602575

Your function TreeIterator._next() is a generator function. This means it returns an iterator upon being called. So you could store the return value of _next() and call .next() on this return value to get the successive elements of this iterator. What you do on the other hand is always return a newly created iterator in TreeIterator.next(), which is never iterated over. This also explains the error message you get: You iterator is not returning tree entries, but rather new iterators.

I think the easiest fix for this problem is to drop the TreeIterator class completely and copy its ._next() method to the .__iter__() method of your tree class. Maybe a few things need to be fixed, but I don't know your tree class.

Upvotes: 1

Nathan Ernst
Nathan Ernst

Reputation: 4590

Looks like you're trying to iterator over the type, not the instance of the type. Change this:

for node in MyTree:
    print "%s: %s" % (node.key(), node.data())

to:

for node in phonebook:
    print "%s: %s" % (node.key(), node.data())

Upvotes: 1

Related Questions