pythonic tree enumerator method using generator

I'm actually using Jython and am pretty new to the Python way of doing things...

When I use javax.swing.tree.DefaultMutableTreeNode I can simply go depth/breadt开发者_运维知识库hFirstEnumeration()...

But if I'm doing stuff with a DOM tree (e.g. from XML) there is no such equivalent... but it strikes me that there must be a very elegant and powerful way of doing this in Python/Jython using a recursive generator.

Hopefully what I want is the most general purpose of utility methods which will essentially do the enumeration with any type of tree object you can throw at it... so you might have to supply the method which gives you the children of a given node... in the case of org.w3c.dom.Node this would be getChildNodes()... then you might want a second optional param which would specify depth or breadth...

To my surprise I haven't been able to find a simple answer just by googling or looking here, for example.

AFAIK, there is no built-in implementation. A very straight-forward solution would be:

import collections

def depth_first_search(node, get_children, depth=0):
    yield node, depth
    for child in get_children(node):
        # In the upcoming Python 3.3, the following can be written as
        # yield from depth_first_search(child, get_children, depth + 1)
        for n, d in depth_first_search(child, get_children, depth + 1):
            yield n, d

def breadth_first_search(node, get_children, depth=0):
    queue = collections.deque([(node, depth)])
    while queue:
        node, depth = queue.popleft()
        queue.extend((n, depth + 1) for n in get_children(node))
        yield node, depth

Then you can easily use these as follows:

def dom_get_children(node):
    nodeList = node.getNodeList()
    for i in range(nodeList.getLength()):
        yield nodeList.item(i)

for node, depth in depth_first_search(some_dom_element, dom_get_children):
    # do something

thanks... while I've been waiting I've been trying to work it out myself..

disclaimer: I wrote this before Ferdinand came up with the definitive version of his excellent answer

in fact your solution, it appears, works fine for a tree consisting of regular Python lists... unfortunately org.w3c.dom.Node is particularly "obtuse"... getChildNodes() actually produces an object called NodeList, which although obviously being a list (Java Array) of some kind remains hermetically sealed from introspection... in particular, dir() will tell you that the class of its "childNodes" field is "org.apache.xerces.dom.DeferredElementImpl"... and my experience is that anything ending in "Impl" ain't gonna be much fun to play with...

I obviously therefore found no way of passing a method as a parameter and invoking it... even with a more amenable Python-like class I'm not currently clear how you invoke a method you pass as a parameter... anyway...

So below are my 3 offerings, pretty self-explanatory: 1) depth-first 2) choice of depth- or breadth-first 3) same but delivering something with an indication of the depth (so you can format print-outs for example). With solution #3 I was forced to create a new class, unfortunately, because I found it was impossible, for example, to add an attribute to the Node object... obviously Jython has limitations and "impurities" compared to Python. I am aware there are python modules for dealing with XML etc... will look into it in due course. (NB of course one of the wonderful aspects of Jython is that you can transition gradually from Java to Python).

Would be interested if any experienced Python/Jython people have any comments...

1. depth-first only:

   def depthFirstTreeEnumeration( node ):
     nodeList = node.getChildNodes()
     for i in range( nodeList.getLength()):
       childNode = nodeList.item( i )
       yield childNode
       for item in depthFirstTreeEnumeration( childNode ):
         yield item
   

2. choice of depth- or breadth-first

   def treeEnumeration( node, depthFirst = True ):
     nodeList = node.getChildNodes()
     for i in range( nodeList.getLength()):
       childNode = nodeList.item( i )
       yield childNode
       if depthFirst:
         for item in treeEnumeration( childNode ):
           yield item
     if not depthFirst:
       for i in range( nodeList.getLength()):
         childNode = nodeList.item( i )
         for item in treeEnumeration( childNode, False ):
           yield item
   

3. choice of depth- or breadth-first, with indication of depth of a given Node

   class NodeWrapper():
     def __init__(self, node, depth ):
       self.node = node
       self.depth = depth
     def __repr__( self ):
       return "node %s, depth %d" % (self.node, self.depth)
   
   def treeEnumerationShowDepth( node, depth = 0, depthFirst = True ):
     nodeList = node.getChildNodes()
     for i in range( nodeList.getLength()):
       wrapper = NodeWrapper( nodeList.item( i ), depth )
       yield wrapper
       if depthFirst:
         for item in treeEnumerationShowDepth( wrapper.node, depth + 1 ):
           yield item
     if not depthFirst:
       for i in range( nodeList.getLength()):
         childNode = nodeList.item( i )
         for item in treeEnumerationShowDepth( childNode, depth + 1, False ):
           yield item
   
   from org.w3c.dom import Node
   
   for wrapper in treeEnumerationShowDepth( dom.getDocumentElement(), 0, False ):
     print "%snode: %s" % ( wrapper.depth * "  ", wrapper.node )