nilmdb/tests/testutil/renderdot.py

import sys

class Renderer(object):

    def __init__(self, getleft, getright,
                 getred, getstart, getend, nil):
        self.getleft = getleft
        self.getright = getright
        self.getred = getred
        self.getstart = getstart
        self.getend = getend
        self.nil = nil

    # Rendering
    def __render_dot_node(self, node, max_depth = 20):
        from nilmdb.utils.printf import sprintf
        """Render a single node and its children into a dot graph fragment"""
        if max_depth == 0:
            return ""
        if node is self.nil:
            return ""
        def c(red):
            if red:
                return 'color="#ff0000", style=filled, fillcolor="#ffc0c0"'
            else:
                return 'color="#000000", style=filled, fillcolor="#c0c0c0"'
        s = sprintf("%d [label=\"%g\\n%g\", %s];\n",
                    id(node),
                    self.getstart(node), self.getend(node),
                    c(self.getred(node)))

        if self.getleft(node) is self.nil:
            s += sprintf("L%d [label=\"-\", %s];\n", id(node), c(False))
            s += sprintf("%d -> L%d [label=L];\n", id(node), id(node))
        else:
            s += sprintf("%d -> %d [label=L];\n",
                         id(node),id(self.getleft(node)))
        if self.getright(node) is self.nil:
            s += sprintf("R%d [label=\"-\", %s];\n", id(node), c(False))
            s += sprintf("%d -> R%d [label=R];\n", id(node), id(node))
        else:
            s += sprintf("%d -> %d [label=R];\n",
                         id(node), id(self.getright(node)))
        s += self.__render_dot_node(self.getleft(node), max_depth-1)
        s += self.__render_dot_node(self.getright(node), max_depth-1)
        return s

    def render_dot(self, rootnode, title = "Tree"):
        """Render the entire tree as a dot graph"""
        return ("digraph rbtree {\n"
                + self.__render_dot_node(rootnode)
                + "}\n");

    def render_dot_live(self, rootnode, title = "Tree"):
        """Render the entiretree as a dot graph, live GTK view"""
        import gtk
        import gtk.gdk
        sys.path.append("/usr/share/xdot")
        import xdot
        xdot.Pen.highlighted = lambda pen: pen
        s = ("digraph rbtree {\n"
             + self.__render_dot_node(rootnode)
             + "}\n");
        window = xdot.DotWindow()
        window.set_dotcode(s)
        window.set_title(title + " - any key to close")
        window.connect('destroy', gtk.main_quit)
        def quit(widget, event):
            if not event.is_modifier:
                window.destroy()
                gtk.main_quit()
        window.widget.connect('key-press-event', quit)
        gtk.main()

class RBTreeRenderer(Renderer):
    def __init__(self, tree):
        Renderer.__init__(self,
                          lambda node: node.left,
                          lambda node: node.right,
                          lambda node: node.red,
                          lambda node: node.start,
                          lambda node: node.end,
                          tree.nil)
        self.tree = tree

    def render(self, title = "RBTree", live = True):
        if live:
            return Renderer.render_dot_live(self, self.tree.getroot(), title)
        else:
            return Renderer.render_dot(self, self.tree.getroot(), title)