diff --git a/migen/graph/treeviz.py b/migen/graph/treeviz.py new file mode 100644 index 000000000..4e4c7750b --- /dev/null +++ b/migen/graph/treeviz.py @@ -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()