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 __future__ import division1718import re1920import blockdiag.elements21from blockdiag.utils import XY2223from rackdiag.utils.math import lcm242526class DiagramNode(blockdiag.elements.DiagramNode):27 pass282930class RackItem(blockdiag.elements.DiagramNode):31 desctable = []32 attrname = {}333435 def clear(cls):36 super(RackItem, cls).clear()37 cls.desctable = ['number', 'label', 'units', 'ampere',38 'weight', 'description']39 cls.attrname = dict(number='No', label='Name', units='Height',40 ampere='Capacity', weight='Weight',41 description='Description')4243 def __init__(self, number, label):44 super(RackItem, self).__init__(None)45 self.number = int(number)46 self.label = label47 self.ampere = None48 self.weight = None4950 if self.label == 'N/A':51 self.color = 'gray'525354 def display_label(self):55 attrs = []56 if self.colheight > 1:57 attrs.append("%dU" % self.colheight)58 if self.ampere:59 attrs.append("%.1fA" % self.ampere)60 if self.weight:61 attrs.append("%.1fkg" % self.weight)6263 labels = []64 if self.label:65 labels.append(self.label)66 if attrs:67 labels.append("[%s]" % "/".join(attrs))6869 return "\n".join(labels)7071 def set_attribute(self, attr):72 if re.search(r'^\d+U$', attr.name):73 self.colheight = int(attr.name[:-1])74 elif re.search(r'^\d+(\.[0-9]+)?A$', attr.name):75 self.ampere = float(attr.name[:-1])76 elif re.search(r'^\d+(\.[0-9]+)?kg$', attr.name):77 self.weight = float(attr.name[:-2])78 else:79 return super(RackItem, self).set_attribute(attr)8081 def set_height(self, value):82 self.colheight = int(value)8384 def to_desctable(self):85 attrs = []86 for name in self.desctable:87 if name == 'units':88 attrs.append("%dU" % self.colheight)89 elif name == 'ampere':90 if self.ampere is None:91 attrs.append("")92 else:93 attrs.append(("%.1fA" % self.ampere) or "")94 elif name == 'weight':95 if self.weight is None:96 attrs.append("")97 else:98 attrs.append(("%.1fkg" % self.weight) or "")99 else:100 value = getattr(self, name)101 if value is None:102 attrs.append("")103 elif isinstance(value, int):104 attrs.append(str(value))105 else:106 attrs.append(getattr(self, name))107108 return attrs109110111class Rack(blockdiag.elements.NodeGroup):112 def __init__(self):113 super(Rack, self).__init__(None)114 self.colheight = 42115 self.description = None116 self.descending = True117118 def set_default_fontsize(self, value):119 raise AttributeError() # Ignore and pass-through to Diagram120121122 def display_label(self):123 attrs = []124 if self.description or attrs:125 attrs.insert(0, "%dU" % self.colheight)126127 labels = []128 if self.description:129 labels.append(self.description)130 if attrs:131 labels.append("[%s]" % "/".join(attrs))132133 return "\n".join(labels)134135 def items(self, levels):136 if isinstance(levels, int):137 levels = [levels]138139 return [node for node in self.nodes if node.xy.y in levels]140141 def get_linked_levels(self, level):142 def get_max_height(self, y):143 try:144 return max(n.colheight for n in self.items(y))145 except Exception:146 return 0147148 height = get_max_height(self, level)149 for dy in range(1, self.colheight - level + 1):150 if dy < height:151 height = max(height, get_max_height(self, level + dy) + dy)152153 return range(level, level + height)154155 def adjust_node_widths(self):156 i = 0157 linked_widths = {}158 widths = []159 while i < self.colheight:160 levels = self.get_linked_levels(i)161162 if levels:163 linked_widths[i] = []164 for level in levels:165 nodes = self.items(level)166 if nodes:167 width = max(n.xy.x for n in nodes) + 1168 widths.append(width)169 linked_widths[i].append(width)170171 i += len(levels)172 else:173 i += 1174175 i = 0176 self.colwidth = lcm(*widths) or 1177 while i < self.colheight:178 levels = self.get_linked_levels(i)179180 if levels:181 colwidth = max(linked_widths[i]) or 1182 for level in levels:183 nodes = self.items(level)184 if nodes:185 width = self.colwidth // colwidth186 for node in nodes:187 node.xy = XY(node.xy.x * width, node.xy.y)188 node.colwidth = width189190 i += len(levels)191 else:192 i += 1193194 def fixiate(self):195 self.adjust_node_widths()196197 for node in self.nodes:198 node.xy = node.xy.shift(x=self.xy.x)199200 def set_ascending(self, attr):201 self.descending = False202203 def set_attribute(self, attr):204 if re.search(r'^\d+U$', attr.name):205 self.colheight = int(attr.name[:-1])206 else:207 return super(Rack, self).set_attribute(attr)208209210class Diagram(blockdiag.elements.Diagram):211 _DiagramNode = RackItem212213214 def clear(cls):215 super(Diagram, cls).clear()216 cls._DiagramNode.clear()217218 def __init__(self):219 super(Diagram, self).__init__()220221 self.racks = [Rack()]222223 def set_default_fontsize(self, value):224 super(Diagram, self).set_default_fontsize(value)225 self.fontsize = int(value)226227 def set_rackheight(self, value):228 self.racks[0].colheight = int(value)229230 def fixiate(self):231 self.colwidth = sum(r.colwidth for r in self.racks)232 self.colheight = max(r.colheight for r in self.racks)233234 def traverse_nodes(self, **kwargs):235 for rack in self.racks:236 for node in rack.traverse_nodes():237 yield node