Jim Paris
f355c73209
There's some bug with the testing harness where placing e.g. from du import du in nilmdb/utils/__init__.py doesn't quite work -- sometimes the module "du" replaces the function "du". Not exactly sure why; we work around that by just renaming files so they don't match the imported names directly.
91 lines
3.3 KiB
Python
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)
|