Draw a sort-of-pretty ascii graph from input data in outline form.

drawgraph.py

#!/usr/bin/env python
import re
import sys
from collections import deque

class Frame:
    def __init__(self,node,indent):
        self.node = node
        self.indent = indent

class PrettyTree:
    """
    Print a tree (actually any directed graph) in pretty ASCII, given an indented outline as input, e.g.
    A
      BB
      CCCC
       DDDDDDDDDDDDDD
       E
      FFF
       CCCC
       GG
    produces:
     _____________A_________
    /         |             \
    BB       CCCC_____     FFF_
             /        \   /    \
       DDDDDDDDDDDDDD E (CCCC) GG
    """

    # Parse outline form input into a dictionary of {node : [children]}.
    def __init__(self,lines,top=None):

        self.graph = {}
        self.nodewidths = {}
        self.top = None
        trail = []

        for line in lines:
            line = line.rstrip()
            m = re.match("([ \.]*)(.*)$",line)  # Space or . indented
            (indent,node) = m.groups()
            indent = len(indent)
            node = re.sub('\s','',node)

            if len(trail)==0:
                trail = [Frame(node,indent)]
                self.graph[node] = []
                self.top = node
                continue

            # more than one contender for root node, so put them all under one we invent
            elif indent==0 or (len(trail)>0 and indent<trail[0].indent):
                trail = [Frame(node,indent)]
                self.graph[node] = []
                if self.top!='(top)':
                    self.graph['(top)'] = [self.top]
                    self.top = '(top)'
                self.graph['(top)'].append(node)
                continue

            name = node

            if  not self.graph.has_key(node):
                self.graph[node] = []

            # if we've encountered this node before, record a
            # childless pointer rather than exploding the tree (or worse)
            if len(self.graph[node])>0:
                name = "(" + name + ")"
                self.graph[name] = []

            # climb back up to the right indentation level
            while indent<=trail[-1].indent:
                trail.pop()

            self.graph[trail[-1].node].append(name)
            trail.append(Frame(node,indent))

    # calculate and memoize the display width of a node, including its descendents.  Without second
    # argument, return whole dictionary of node widths.
    def nodewidth(self,node=None):
        if node==None:
            self.nodewidth(self.top)
            return self.nodewidths
        if self.nodewidths.has_key(node):
            return self.nodewidths[node]
        # Add up width of children, including space delimiters
        w = len(self.graph[node])-1 if len(self.graph[node])>0 else 0
        for child in self.graph[node]:
            w = w + self.nodewidth(child)
        w = len(node) if w<len(node) else w # unless parent's name is itself really, really long
        self.nodewidths[node] = w
        return w

    # pre and post-pad node names with underscore, to connect to bars below
    def padnodes(self,line,bars):
        leftyet = False
        ret = ''
        for (l,b) in zip(list(line),list(bars)):
            if b=='/':
                leftyet = True
            elif b=='\\':
                leftyet = False
            elif leftyet and l==' ':
                l = '_'
            ret += l
        if len(line)>len(bars):
            ret += line[len(bars):]
        return ret

    def layout(self):
        queue = deque([(self.top,0,0,'|')])  # node, parent indent, depth, bar pointing to us
        lines = ''  # output
        line = ''   # current line
        pline = ''  # parent line
        bars = ''   # the bars (/,\,|) pointing to nodes on current line
        pdepth = 0
        # BFS traversal over graph.  At each depth, accumulate the text line to print out.
        # Since each line must be jiggered based on the position and length of node names on
        # the following line, this is a bit complicated.
        while len(queue)>0:
            (node,pind,depth,bar) = queue.popleft() # node, parent indent, depth, bar pointing to us
            assert depth>=pdepth, 'Horrible error ' + str((node,depth,pdepth))
            # If we've stepped down to the next level, print out the level above and the bars
            if depth>pdepth:
                pdepth = depth
                if pdepth>1:
                    lines += self.padnodes(pline,bars) + '\n'
                    lines += bars + '\n'
                pline = line
                line = ''
                bars = ''
            # Indent at least to indentation of parent node
            line += (pind-len(line))*' '
            bars += (pind-len(bars))*' '
            n = len(self.graph[node])
            # enqueue children, figuring proper sort of bar to point to them
            for (i,c) in enumerate(self.graph[node]):
                cbar = '/' if i==0 and n>1 else '\\' if i==n-1 and n>1 else '|'
                queue.append((c,len(line),depth+1,cbar))

            # Pre and post-pad node names so they're centered over their descendents
            nl = len(node)
            w = self.nodewidth(node)
            wl = (w-nl)/2
            wr = w - wl - nl+1
            line += wl*' ' + node + wr*' '
            bars += (wl+(nl-1)/2)*' ' + bar + (wr+(nl-(nl-1)/2-1))*' '

        lines += self.padnodes(pline,bars) + '\n'
        lines += bars + '\n' + line
        return lines

    def __str__(self):
        return self.layout()


if __name__ == "__main__":
    lines = sys.stdin.readlines()
    T = PrettyTree(lines)
    print T.layout()