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 namedtuple1718from blockdiag.utils import XY, unquote1920from nwdiag import parser21from nwdiag.elements import (Diagram, DiagramEdge, DiagramNode, Network,22 NodeGroup, Route)232425class DiagramTreeBuilder:26 def build(self, tree):27 self.diagram = Diagram()28 self.instantiate(None, None, tree)29 for network in self.diagram.networks:30 nodes = [n for n in self.diagram.nodes if network in n.networks]31 if len(nodes) == 0:32 self.diagram.networks.remove(network)3334 for i, network in enumerate(self.diagram.networks):35 network.xy = XY(0, i)3637 for subgroup in self.diagram.groups:38 if len(subgroup.nodes) == 0:39 self.diagram.groups.remove(subgroup)4041 for node in self.diagram.nodes:42 if len(node.networks) == 0:43 msg = "DiagramNode %s does not belong to any networks"44 raise RuntimeError(msg % node.id)4546 # show networks including same nodes47 for nw in self.diagram.networks:48 if nw.hidden and len(nw.nodes) == 2:49 def is_same(x):50 return set(nodes) & set(x.nodes) == set(nodes)5152 nodes = nw.nodes53 for n in self.diagram.networks:54 if n != nw and is_same(n):55 nw.hidden = False56 break5758 # show network for multiple peer networks from same node59 nodes = []60 for nw in self.diagram.networks:61 if nw.hidden and len(nw.nodes) == 2:62 nodes.append(nw.nodes[0]) # parent node (FROM node)6364 for node in [n for n in set(nodes) if nodes.count(n) > 1]:65 for network in node.networks:66 if len(network.nodes) == 2:67 network.hidden = False6869 return self.diagram7071 def instantiate(self, network, group, tree):72 for stmt in tree.stmts:73 if isinstance(stmt, parser.Node):74 node = DiagramNode.get(stmt.id)75 node.set_attributes(network, stmt.attrs)7677 if group:78 if (node.group and node.group != self.diagram and79 (group != node.group)):80 msg = "DiagramNode could not belong to two groups"81 raise RuntimeError(msg)82 node.group = group83 group.nodes.append(node)84 elif node.group is None:85 node.group = self.diagram8687 if network is not None:88 if network not in node.networks:89 node.networks.append(network)90 if node not in network.nodes:91 network.nodes.append(node)92 if node not in self.diagram.nodes:93 self.diagram.nodes.append(node)9495 elif isinstance(stmt, parser.Network):96 subnetwork = Network.get(stmt.id)97 subnetwork.label = stmt.id9899 if subnetwork not in self.diagram.networks:100 self.diagram.networks.append(subnetwork)101102 substmt = namedtuple('Statements', 'stmts')([])103 for s in stmt.stmts:104 if isinstance(s, parser.Attr):105 subnetwork.set_attribute(s)106 else:107 substmt.stmts.append(s)108109 self.instantiate(subnetwork, group, substmt)110111 elif isinstance(stmt, parser.Group):112 subgroup = NodeGroup.get(stmt.id)113114 if subgroup not in self.diagram.groups:115 self.diagram.groups.append(subgroup)116117 substmt = namedtuple('Statements', 'stmts')([])118 for s in stmt.stmts:119 if isinstance(s, parser.Attr):120 subgroup.set_attribute(s)121 else:122 substmt.stmts.append(s)123124 self.instantiate(network, subgroup, substmt)125126 elif isinstance(stmt, parser.Peer):127 nodes = []128 for edge in stmt.edges:129 nodes.append(DiagramNode.get(edge.from_node))130 nodes.append(DiagramNode.get(edge.to_node))131132 for node in nodes:133 if node.group is None:134 node.group = self.diagram135 if node not in self.diagram.nodes:136 self.diagram.nodes.append(node)137138 if len(nodes[0].networks) == 0:139 nw = Network.create_anonymous([nodes[0]])140 if nw:141 self.diagram.networks.append(nw)142143 for edge in stmt.edges:144 from_node = DiagramNode.get(edge.from_node)145 to_node = DiagramNode.get(edge.to_node)146147 _nodes = [from_node, to_node]148 nw = Network.create_anonymous(_nodes, edge.attrs)149 if nw:150 self.diagram.networks.append(nw)151152 elif isinstance(stmt, parser.Route):153 nodes = [DiagramNode.get(n) for n in stmt.nodes]154 for node1, node2 in zip(nodes[:-1], nodes[1:]):155 route = Route(node1, node2)156 route.set_attributes(stmt.attrs)157 self.diagram.routes.append(route)158159 elif isinstance(stmt, parser.Attr):160 self.diagram.set_attribute(stmt)161162 elif isinstance(stmt, parser.Extension):163 if stmt.type == 'class':164 name = unquote(stmt.name)165 Diagram.classes[name] = stmt166 elif stmt.type == 'plugin':167 self.diagram.set_plugin(stmt.name, stmt.attrs)168169 elif isinstance(stmt, parser.Statements):170 self.instantiate(network, group, stmt)171172 return network173174175class DiagramLayoutManager:176 def __init__(self, diagram):177 self.diagram = diagram178179 self.coordinates = []180181 def run(self):182 self.do_layout()183 self.diagram.fixiate()184185 def do_layout(self):186 self.layout_nodes()187 self.set_network_size()188189 def layout_nodes(self, group=None):190 networks = self.diagram.networks191192 if group:193 nodes = (n for n in self.diagram.nodes if n.group == group)194 else:195 nodes = self.diagram.nodes196197 for node in nodes:198 if node.layouted:199 continue200201 joined = [g for g in node.networks if g.hidden is False]202 y1 = min(networks.index(g) for g in node.networks)203 if joined:204 y2 = max(networks.index(g) for g in joined)205 else:206 y2 = y1207208 if node.group and node.group != self.diagram and group:209 starts = min(n.xy.x for n in group.nodes if n.layouted)210 else:211 nw = [n for n in node.networks if n.xy.y == y1][0]212 nodes = [n for n in self.diagram.nodes if nw in n.networks]213 layouted = [n for n in nodes if n.xy.x > 0]214215 starts = 0216 if layouted:217 layouted.sort(key=lambda a: a.xy.x)218 basenode = layouted[0]219 commonnw = set(basenode.networks) & set(node.networks)220221 if basenode.xy.y == y1:222 starts = basenode.xy.x + 1223 elif commonnw and list(commonnw)[0].hidden is True:224 starts = basenode.xy.x225 else:226 starts = basenode.xy.x + 1 - len(nodes)227228 if starts < 0:229 starts = 0230231 for x in range(starts, len(self.diagram.nodes)):232 points = [XY(x, y) for y in range(y1, y2 + 1)]233 if not set(points) & set(self.coordinates):234 node.xy = XY(x, y1)235 node.layouted = True236 self.coordinates += points237 break238239 if node.group and node.group != self.diagram and group is None:240 self.layout_nodes(node.group)241242 if group:243 self.set_coordinates(group)244245 def set_coordinates(self, group):246 self.set_group_size(group)247248 xy = group.xy249 for i in range(xy.x, xy.x + group.colwidth):250 for j in range(xy.y, xy.y + group.colheight):251 self.coordinates.append(XY(i, j))252253 def set_network_size(self):254 for network in self.diagram.networks:255 nodes = [n for n in self.diagram.nodes if network in n.networks]256 nodes.sort(key=lambda a: a.xy.x)257258 x0 = min(n.xy.x for n in nodes)259 network.xy = XY(x0, network.xy.y)260261 x1 = max(n.xy.x for n in nodes)262 network.colwidth = x1 - x0 + 1263264 def set_group_size(self, group):265 nodes = list(group.nodes)266 nodes.sort(key=lambda a: a.xy.x)267268 x0 = min(n.xy.x for n in nodes)269 y0 = min(n.xy.y for n in nodes)270 group.xy = XY(x0, y0)271272 x1 = max(n.xy.x for n in nodes)273 y1 = max(n.xy.y for n in nodes)274 group.colwidth = x1 - x0 + 1275 group.colheight = y1 - y0 + 1276277278class ScreenNodeBuilder:279280 def build(cls, tree):281 DiagramNode.clear()282 DiagramEdge.clear()283 NodeGroup.clear()284 Network.clear()285286 diagram = DiagramTreeBuilder().build(tree)287 DiagramLayoutManager(diagram).run()288 diagram = cls.update_network_status(diagram)289290 return diagram291292293 def update_network_status(cls, diagram):294 for node in diagram.nodes:295 above = [nw for nw in node.networks if nw.xy.y <= node.xy.y]296 if len(above) > 1 and [nw for nw in above if nw.hidden]:297 for nw in above:298 nw.hidden = False299300 below = [nw for nw in node.networks if nw.xy.y > node.xy.y]301 if len(below) > 1 and [nw for nw in below if nw.hidden]:302 for nw in below:303 nw.hidden = False304305 return diagram