add tree visualizer
authorSebastien Bourdeauducq <sebastien@milkymist.org>
Wed, 7 Aug 2013 13:52:35 +0000 (15:52 +0200)
committerSebastien Bourdeauducq <sebastien@milkymist.org>
Wed, 7 Aug 2013 13:52:35 +0000 (15:52 +0200)
migen/graph/treeviz.py [new file with mode: 0644]

diff --git a/migen/graph/treeviz.py b/migen/graph/treeviz.py
new file mode 100644 (file)
index 0000000..4e4c775
--- /dev/null
@@ -0,0 +1,88 @@
+import cairo
+import math
+
+def _cairo_draw_node(ctx, radius, color, outer_color, s):
+       ctx.save()
+
+       ctx.set_line_width(0.0)
+       gradient_color = cairo.RadialGradient(0, 0, 0, 0, 0, radius)
+       gradient_color.add_color_stop_rgb(0, *color)
+       gradient_color.add_color_stop_rgb(1, *outer_color)
+       ctx.set_source(gradient_color)
+       ctx.arc(0, 0, radius, 0, 2*math.pi)
+       ctx.fill()
+
+       ctx.set_source_rgb(0, 0, 0)
+       x_bearing, y_bearing, textw, texth, x_advance, y_advance = ctx.text_extents(s)
+       ctx.translate(-textw/2, texth/2)
+       ctx.show_text(s)
+
+       ctx.restore()
+
+def _cairo_draw_connection(ctx, x0, y0, color0, x1, y1, color1):
+       ctx.move_to(x0, y0)
+       ctx.curve_to(x0, y0+20, x1, y1-20, x1, y1)
+       ctx.set_line_width(1.2)
+       gradient_color = cairo.LinearGradient(x0, y0, x1, y1)
+       gradient_color.add_color_stop_rgb(0, *color0)
+       gradient_color.add_color_stop_rgb(1, *color1)
+       ctx.set_source(gradient_color)
+       ctx.stroke()
+
+class RenderNode:
+       def __init__(self, label, children=None, color=(0.8, 0.8, 0.8), radius=40):
+               self.label = label
+               if children is None:
+                       children = []
+               self.children = children
+               self.color = color
+               self.outer_color = (color[0]*3/5, color[1]*3/5, color[2]*3/5)
+               self.radius = radius
+               self.pitch = self.radius*3
+
+       def get_extents(self):
+               if self.children:
+                       cw, ch = zip(*[c.get_extents() for c in self.children])
+                       w = max(cw)*len(self.children)
+                       h = self.pitch + max(ch)
+               else:
+                       w = h = self.pitch
+               return w, h
+
+       def render(self, ctx):
+               _cairo_draw_node(ctx, self.radius, self.color, self.outer_color, self.label)
+               if self.children:
+                       cpitch = max([c.get_extents()[0] for c in self.children])
+                       first_child_x = -(cpitch*(len(self.children) - 1))/2
+
+                       ctx.save()
+                       ctx.translate(first_child_x, self.pitch)
+                       for c in self.children:
+                               c.render(ctx)
+                               ctx.translate(cpitch, 0)
+                       ctx.restore()
+
+                       current_x = first_child_x
+                       for c in self.children:
+                               current_y = self.pitch - c.radius
+                               _cairo_draw_connection(ctx, 0, self.radius, self.outer_color, current_x, current_y, c.outer_color)
+                               current_x += cpitch
+
+       def to_svg(self, name):
+               w, h = self.get_extents()
+               surface = cairo.SVGSurface(name, w, h)
+               ctx = cairo.Context(surface)
+               ctx.translate(w/2, self.pitch/2)
+               self.render(ctx)
+               surface.finish()
+
+def _test():
+       xns = [RenderNode("X"+str(n)) for n in range(5)]
+       yns = [RenderNode("Y"+str(n), [RenderNode("foo", color=(0.1*n, 0.5+0.2*n, 1.0-0.3*n))]) for n in range(3)]
+       n1 = RenderNode("n1", yns)
+       n2 = RenderNode("n2", xns, color=(0.8, 0.5, 0.9))
+       top = RenderNode("top", [n1, n2])
+       top.to_svg("test.svg")
+
+if __name__ == "__main__":
+       _test()