migen/util/treeviz.py

   1 import cairo
   2 import math
   3
   4
   5 def _cairo_draw_node(ctx, dx, radius, color, outer_color, s):
   6     ctx.save()
   7
   8     ctx.translate(dx, 0)
   9
  10     ctx.set_line_width(0.0)
  11     gradient_color = cairo.RadialGradient(0, 0, 0, 0, 0, radius)
  12     gradient_color.add_color_stop_rgb(0, *color)
  13     gradient_color.add_color_stop_rgb(1, *outer_color)
  14     ctx.set_source(gradient_color)
  15     ctx.arc(0, 0, radius, 0, 2*math.pi)
  16     ctx.fill()
  17
  18     lines = s.split("\n")
  19     textws = []
  20     texths = []
  21     for line in lines:
  22         x_bearing, y_bearing, w, h, x_advance, y_advance = ctx.text_extents(line)
  23         textws.append(w)
  24         texths.append(h + 2)
  25     ctx.translate(0, -sum(texths[1:])/2)
  26     for line, w, h in zip(lines, textws, texths):
  27         ctx.translate(-w/2, h/2)
  28         ctx.move_to(0, 0)
  29         ctx.set_source_rgb(0, 0, 0)
  30         ctx.show_text(line)
  31         ctx.translate(w/2, h/2)
  32
  33     ctx.restore()
  34
  35
  36 def _cairo_draw_connection(ctx, x0, y0, color0, x1, y1, color1):
  37     ctx.move_to(x0, y0)
  38     ctx.curve_to(x0, y0+20, x1, y1-20, x1, y1)
  39     ctx.set_line_width(1.2)
  40     gradient_color = cairo.LinearGradient(x0, y0, x1, y1)
  41     gradient_color.add_color_stop_rgb(0, *color0)
  42     gradient_color.add_color_stop_rgb(1, *color1)
  43     ctx.set_source(gradient_color)
  44     ctx.stroke()
  45
  46
  47 class RenderNode:
  48     def __init__(self, label, children=None, color=(0.8, 0.8, 0.8), radius=40):
  49         self.label = label
  50         if children is None:
  51             children = []
  52         self.children = children
  53         self.color = color
  54         self.outer_color = (color[0]*3/5, color[1]*3/5, color[2]*3/5)
  55         self.radius = radius
  56         self.pitch = self.radius*3
  57
  58     def get_dimensions(self):
  59         if self.children:
  60             cws, chs, cdxs = zip(*[c.get_dimensions() for c in self.children])
  61             w = sum(cws)
  62             h = self.pitch + max(chs)
  63             dx = cws[0]/4 - cws[-1]/4
  64         else:
  65             w = h = self.pitch
  66             dx = 0
  67         return w, h, dx
  68
  69     def render(self, ctx):
  70         if self.children:
  71             cws, chs, cdxs = zip(*[c.get_dimensions() for c in self.children])
  72             first_child_x = -sum(cws)/2
  73
  74             ctx.save()
  75             ctx.translate(first_child_x, self.pitch)
  76             for c, w in zip(self.children, cws):
  77                 ctx.translate(w/2, 0)
  78                 c.render(ctx)
  79                 ctx.translate(w/2, 0)
  80             ctx.restore()
  81
  82             dx = cws[0]/4 - cws[-1]/4
  83
  84             current_x = first_child_x
  85             for c, w, cdx in zip(self.children, cws, cdxs):
  86                 current_y = self.pitch - c.radius
  87                 current_x += w/2
  88                 _cairo_draw_connection(ctx, dx, self.radius, self.outer_color, current_x+cdx, current_y, c.outer_color)
  89                 current_x += w/2
  90         else:
  91             dx = 0
  92         _cairo_draw_node(ctx, dx, self.radius, self.color, self.outer_color, self.label)
  93
  94     def to_svg(self, name):
  95         w, h, dx = self.get_dimensions()
  96         surface = cairo.SVGSurface(name, w, h)
  97         ctx = cairo.Context(surface)
  98         ctx.translate(w/2, self.pitch/2)
  99         self.render(ctx)
 100         surface.finish()
 101
 102
 103 def _test():
 104     xns = [RenderNode("X"+str(n)) for n in range(5)]
 105     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)]
 106     n1 = RenderNode("n1", yns)
 107     n2 = RenderNode("n2", xns, color=(0.8, 0.5, 0.9))
 108     top = RenderNode("top", [n1, n2])
 109     top.to_svg("test.svg")
 110
 111 if __name__ == "__main__":
 112     _test()