nilmdb/tests/testutil/renderdot.py

91 lines
3.3 KiB
Python

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)