1# -*- coding: utf-8 -*-2# Copyright 2011 Takeshi KOMIYA3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.1516from collections import defaultdict1718from blockdiag import imagedraw, noderenderer19from blockdiag.metrics import AutoScaler, DiagramMetrics202122class DiagramDraw(object):23 shadow_colors = defaultdict(lambda: (0, 0, 0))24 shadow_colors['PNG'] = (64, 64, 64)25 shadow_colors['PDF'] = (144, 144, 144)2627 def __init__(self, _format, diagram, filename=None, **kwargs):28 self.format = _format.upper()29 self.diagram = diagram30 self.fill = kwargs.get('fill', (0, 0, 0))31 self.badgeFill = kwargs.get('badgeFill', 'pink')32 self.filename = filename33 self.shadow = self.shadow_colors[self.format.upper()]3435 if self.format == 'PNG' and kwargs.get('antialias'):36 self.scale_ratio = 237 else:38 self.scale_ratio = 13940 self.drawer = imagedraw.create(self.format, self.filename,41 filters=['linejump'],42 scale_ratio=self.scale_ratio,43 **kwargs)4445 self.metrics = self.create_metrics(kwargs.get('basediagram', diagram),46 drawer=self.drawer,47 fontmap=kwargs.get('fontmap'))48 if self.scale_ratio == 2:49 self.metrics = AutoScaler(self.metrics,50 scale_ratio=self.scale_ratio)5152 self.drawer.set_canvas_size(self.pagesize())53 self.drawer.set_options(jump_radius=self.metrics.cellsize / 2)5455 def create_metrics(self, *args, **kwargs):56 return DiagramMetrics(*args, **kwargs)575859 def nodes(self):60 for node in self.diagram.traverse_nodes():61 if node.drawable:62 yield node636465 def groups(self):66 for group in self.diagram.traverse_groups(preorder=True):67 if not group.drawable:68 yield group697071 def edges(self):72 edges = self.diagram.traverse_edges(preorder=True)73 for edge in (e for e in edges if e.style != 'none'):74 yield edge7576 def pagesize(self, scaled=False):77 if scaled:78 metrics = self.metrics79 else:80 metrics = self.metrics.original_metrics8182 width = self.diagram.colwidth83 height = self.diagram.colheight84 return metrics.pagesize(width, height)8586 def draw(self, **kwargs):87 # switch metrics object during draw backgrounds88 temp, self.metrics = self.metrics, self.metrics.original_metrics89 self._draw_background()90 self.metrics = temp9192 if self.scale_ratio > 1:93 pagesize = self.pagesize(scaled=True)94 self.drawer.resizeCanvas(pagesize)9596 self._draw_elements(**kwargs)9798 def _draw_background(self):99 # Draw node groups.100 for group in self.groups:101 if group.shape == 'box':102 box = self.metrics.group(group).marginbox103 if group.href and self.format == 'SVG':104 drawer = self.drawer.anchor(group.href)105 else:106 drawer = self.drawer107108 drawer.rectangle(box, fill=group.color, filter='blur')109110 # Drop node shadows.111 for node in self.nodes:112 if node.color != 'none' and self.diagram.shadow_style != 'none':113 r = noderenderer.get(node.shape)114115 shape = r(node, self.metrics)116 if node.href and self.format == 'SVG':117 drawer = self.drawer.anchor(node.href)118 else:119 drawer = self.drawer120121 shape.render(drawer, self.format,122 fill=self.shadow, shadow=True,123 style=self.diagram.shadow_style)124125 def _draw_elements(self, **kwargs):126 for node in self.nodes:127 self.node(node, **kwargs)128129 for edge in self.edges:130 self.edge(edge)131132 for edge in self.edges:133 self.edge_label(edge)134135 for group in self.groups:136 if group.shape == 'line':137 box = self.metrics.group(group).marginbox138 self.drawer.rectangle(box, fill='none', outline=group.color,139 style=group.style, thick=group.thick)140141 for node in self.groups:142 self.group_label(node, **kwargs)143144 def node(self, node, **kwargs):145 r = noderenderer.get(node.shape)146 shape = r(node, self.metrics)147 if node.href and self.format == 'SVG':148 drawer = self.drawer.anchor(node.href)149 else:150 drawer = self.drawer151152 shape.render(drawer, self.format, fill=self.fill,153 badgeFill=self.badgeFill)154155 def group_label(self, group):156 m = self.metrics.group(group)157 font = self.metrics.font_for(group)158159 if group.label and not group.separated:160 self.drawer.textarea(m.grouplabelbox, group.label, font=font,161 fill=group.textcolor)162 elif group.label:163 self.drawer.textarea(m.corebox, group.label, font=font,164 fill=group.textcolor)165166 def edge(self, edge):167 metrics = self.metrics.edge(edge)168169 for line in metrics.shaft.polylines:170 self.drawer.line(line, fill=edge.color, thick=edge.thick,171 style=edge.style, jump=True)172173 for head in metrics.heads:174 if edge.hstyle in ('generalization', 'aggregation'):175 self.drawer.polygon(head, outline=edge.color, fill='white')176 else:177 self.drawer.polygon(head, outline=edge.color, fill=edge.color)178179 def edge_label(self, edge):180 if edge.label:181 metrics = self.metrics.edge(edge)182183 font = self.metrics.font_for(edge)184 self.drawer.textarea(metrics.labelbox, edge.label, font=font,185 fill=edge.textcolor, outline=self.fill)186187 def save(self, size=None):188 return self.drawer.save(self.filename, size, self.format)