Getting a rooted graph from a nested list of lists

edit

One way to do this would be to modify the original tree function to parse the list; another would be to build the graph directly. The second one is slightly trickier because the names of graph vertices must be unique, so you have to ensure this with some kind of global counter. Here's one version I made by writing a class to contain both the counter, the graph, and the recursive tree builder:

class ExpressionGraph():
    def __init__(self,expr):
        self.G = Graph()
        self.i = 0
        self.expr = expr
        self.root = None
        self.graph_expr(self.expr)
    def plot(self,*args,**kwds):
        #print "root is {0}".format(self.root)
        return self.G.plot(*args,layout='tree',tree_root=self.root,**kwds)
    def graph_expr(self,expr):
        try:
            operator = expr.operator()
        except AttributeError: # e.g. if expr is an integer
            operator = None

        if operator is None:
            name = "[{0}] {1}".format(self.i,expr)
            print self.i
            print "(leaf) {0}".format(expr)
            self.i += 1
            self.G.add_vertex(name)
            return name
        else:
            name = "[{0}] {1}".format(self.i,expr.operator().__name__)
            print self.i
            print "(node) {0}; {1}".format(expr,expr.operator().__name__)
            if self.i == 0:
                self.root = name
                print "  ** root is '{0}' **".format(self.root)
            self.i += 1
            new_nodes = []
            for opnd in expr.operands():
                new_nodes += [self.graph_expr(opnd)]
            self.G.add_vertex(name)
            self.G.add_edges([(name,node) for node in new_nodes])
            return name

Note that I've included some print statements so you can see how it's working:

f = 2*(x+2^x)
E = ExpressionGraph(f)

prints out

0
(node) 2*2^x + 2*x; add
  ** root is '[0] add' **
1
(node) 2*2^x; mul
2
(node) 2^x; pow
3
(leaf) 2
4
(leaf) x
5
(leaf) 2
6
(node) 2*x; mul
7
(leaf) x
8
(leaf) 2

and E.plot() plots the tree, with additional arguments passed to the underlying graph. For example:

E.plot(vertex_size=2000,vertex_colors='#ffccff')