Can you define multiple different iterators for a Python class?

Question

I'm writing a very simple Tree class:

class Tree:
    def __init__(self, value_ = None, children_ = None):
        self.value = value_
        self.children = children_

I'd like to be able to perform both DFS and BFS traversal with a simple loop, i.e.:

t = Tree()
# ...fill tree...

for node in t:
    print(node.value)

In C++, for example, you can have multiple types of iterators - so I could define both a DFS and a BFS iterator and use one or the other depending on what type of traversal I wanted to do. Is this possible to do in Python?

Answer 1

You can have multiple methods returning iterators and have the 'default' one as __iter__. Below is a simple binary tree where 'default' iterator does DFS and which additionally supports BFS with separate method:

from collections import deque

class Tree(object):
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

    def __iter__(self):
        if self.left:
            for x in self.left:
                yield x

        yield self.value

        if self.right:
            for x in self.right:
                yield x

    def bfs(self):
        q = deque([self])
        while q:
            x = q.popleft()
            if x:
                yield x.value
                q.extend([x.left, x.right])

Short example of usage:

root = Tree(2)
root.left = Tree(1)
root.right = Tree(4)
root.right.left = Tree(3)
root.right.right = Tree(5)

print list(root) # [1, 2, 3, 4, 5]
print list(root.bfs()) # [2, 1, 4, 3, 5]

Answer 2

You could write separate methods on your class for the two types of iteration. These could for instance be generators that yield values in whatever order you want. You would then write something like:

for node in t.depth_first():
    # ...

for node in t.breadth_first():
    # ...