title: Implementing a depth-first tree iterator in Python
body: <p>I'm trying to implement an iterator class for not-necessarily-binary trees in Python. After the iterator is constructed with a tree's root node, its <code>next()</code> function can be called repeatedly to traverse the tree in depth-first order (e.g., <a href="http://en.wikipedia.org/wiki/Depth-first_search#mediaviewer/File:Depth-first-tree.svg" rel="noreferrer">this order</a>), finally returning <code>None</code> when there are no nodes left. </p>
  
  <p>Here is the basic <code>Node</code> class for a tree:</p>
  
  <pre><code>class Node(object):
  
      def __init__(self, title, children=None):
          self.title = title
          self.children = children or []
          self.visited = False   
  
      def __str__(self):
          return self.title
  </code></pre>
  
  <p>As you can see above, I introduced a <code>visited</code> property to the nodes for my first approach, since I didn't see a way around it. With that extra measure of state, the <code>Iterator</code> class looks like this:</p>
  
  <pre><code>class Iterator(object):
  
      def __init__(self, root):
          self.stack = []
          self.current = root
  
      def next(self):
          if self.current is None:
              return None
  
          self.stack.append(self.current)
          self.current.visited = True
  
          # Root case
          if len(self.stack) == 1:
              return self.current
  
          while self.stack:
              self.current = self.stack[-1] 
              for child in self.current.children:
                  if not child.visited:
                      self.current = child
                      return child
  
              self.stack.pop()
  </code></pre>
  
  <p>This is all well and good, but I want to get rid of the need for the <code>visited</code> property, without resorting to recursion or any other alterations to the <code>Node</code> class. </p>
  
  <p>All the state I need should be taken care of in the iterator, but I'm at a loss about how that can be done. Keeping a visited list for the whole tree is non-scalable and out of the question, so there must be a clever way to use the stack. </p>
  
  <p>What especially confuses me is this--since the <code>next()</code> function, of course, returns, how can I remember where I've been without marking anything or using excess storage? Intuitively, I think of looping over children, but that logic is broken/forgotten when the <code>next()</code> function returns! </p>
  
  <p>UPDATE - Here is a small test:</p>
  
  <pre><code>tree = Node(
      'A', [
          Node('B', [
              Node('C', [
                  Node('D')
                  ]),
              Node('E'),
              ]),
          Node('F'),
          Node('G'),
          ])
  
  iter = Iterator(tree)
  
  out = object()
  while out:
      out = iter.next()
      print out
  </code></pre>