1# -*- coding: utf-8 -*-2import sys34from graphviz import Graph5from collections import Counter6from typing import Any, List, Union7from dataclasses import dataclass8from pathlib import Path9from itertools import zip_longest10import re1112from wireviz import wv_colors, __version__, APP_NAME, APP_URL13from wireviz.DataClasses import Metadata, Options, Tweak, Connector, Cable14from wireviz.wv_colors import get_color_hex, translate_color15from wireviz.wv_gv_html import nested_html_table, \16 html_bgcolor_attr, html_bgcolor, html_colorbar, \17 html_image, html_caption, remove_links, html_line_breaks18from wireviz.wv_bom import pn_info_string, component_table_entry, \19 get_additional_component_table, bom_list, generate_bom, \20 HEADER_PN, HEADER_MPN, HEADER_SPN21from wireviz.wv_html import generate_html_output22from wireviz.wv_helper import awg_equiv, mm2_equiv, tuplelist2tsv, flatten2d, \23 open_file_read, open_file_write24252627class Harness:28 metadata: Metadata29 options: Options30 tweak: Tweak3132 def __post_init__(self):33 self.connectors = {}34 self.cables = {}35 self._bom = [] # Internal Cache for generated bom36 self.additional_bom_items = []3738 def add_connector(self, name: str, *args, **kwargs) -> None:39 self.connectors[name] = Connector(name, *args, **kwargs)4041 def add_cable(self, name: str, *args, **kwargs) -> None:42 self.cables[name] = Cable(name, *args, **kwargs)4344 def add_bom_item(self, item: dict) -> None:45 self.additional_bom_items.append(item)4647 def connect(self, from_name: str, from_pin: (int, str), via_name: str, via_wire: (int, str), to_name: str, to_pin: (int, str)) -> None:48 # check from and to connectors49 for (name, pin) in zip([from_name, to_name], [from_pin, to_pin]):50 if name is not None and name in self.connectors:51 connector = self.connectors[name]52 # check if provided name is ambiguous53 if pin in connector.pins and pin in connector.pinlabels:54 if connector.pins.index(pin) != connector.pinlabels.index(pin):55 raise Exception(f'{name}:{pin} is defined both in pinlabels and pins, for different pins.')56 # TODO: Maybe issue a warning if present in both lists but referencing the same pin?57 if pin in connector.pinlabels:58 if connector.pinlabels.count(pin) > 1:59 raise Exception(f'{name}:{pin} is defined more than once.')60 index = connector.pinlabels.index(pin)61 pin = connector.pins[index] # map pin name to pin number62 if name == from_name:63 from_pin = pin64 if name == to_name:65 to_pin = pin66 if not pin in connector.pins:67 raise Exception(f'{name}:{pin} not found.')6869 # check via cable70 if via_name in self.cables:71 cable = self.cables[via_name]72 # check if provided name is ambiguous73 if via_wire in cable.colors and via_wire in cable.wirelabels:74 if cable.colors.index(via_wire) != cable.wirelabels.index(via_wire):75 raise Exception(f'{via_name}:{via_wire} is defined both in colors and wirelabels, for different wires.')76 # TODO: Maybe issue a warning if present in both lists but referencing the same wire?77 if via_wire in cable.colors:78 if cable.colors.count(via_wire) > 1:79 raise Exception(f'{via_name}:{via_wire} is used for more than one wire.')80 via_wire = cable.colors.index(via_wire) + 1 # list index starts at 0, wire IDs start at 181 elif via_wire in cable.wirelabels:82 if cable.wirelabels.count(via_wire) > 1:83 raise Exception(f'{via_name}:{via_wire} is used for more than one wire.')84 via_wire = cable.wirelabels.index(via_wire) + 1 # list index starts at 0, wire IDs start at 18586 from_pin_id = self.connectors[from_name].pins.index(from_pin) if from_pin is not None else None87 to_pin_id = self.connectors[to_name].pins.index(to_pin) if to_pin is not None else None8889 self.cables[via_name].connect(from_name, from_pin_id, via_wire, to_name, to_pin_id)90 if from_name in self.connectors:91 self.connectors[from_name].activate_pin(from_pin)92 if to_name in self.connectors:93 self.connectors[to_name].activate_pin(to_pin)9495 def create_graph(self) -> Graph:96 dot = Graph()97 dot.body.append(f'// Graph generated by {APP_NAME} {__version__}\n')98 dot.body.append(f'// {APP_URL}\n')99 dot.attr('graph', rankdir='LR',100 ranksep='2',101 bgcolor=wv_colors.translate_color(self.options.bgcolor, "HEX"),102 nodesep='0.33',103 fontname=self.options.fontname)104 dot.attr('node',105 shape='none',106 width='0', height='0', margin='0', # Actual size of the node is entirely determined by the label.107 style='filled',108 fillcolor=wv_colors.translate_color(self.options.bgcolor_node, "HEX"),109 fontname=self.options.fontname)110 dot.attr('edge', style='bold',111 fontname=self.options.fontname)112113 # prepare ports on connectors depending on which side they will connect114 for _, cable in self.cables.items():115 for connection_color in cable.connections:116 if connection_color.from_port is not None: # connect to left117 self.connectors[connection_color.from_name].ports_right = True118 if connection_color.to_port is not None: # connect to right119 self.connectors[connection_color.to_name].ports_left = True120121 for connector in self.connectors.values():122123 # If no wires connected (except maybe loop wires)?124 if not (connector.ports_left or connector.ports_right):125 connector.ports_left = True # Use left side pins.126127 html = []128129 rows = [[f'{html_bgcolor(connector.bgcolor_title)}{remove_links(connector.name)}'130 if connector.show_name else None],131 [pn_info_string(HEADER_PN, None, remove_links(connector.pn)),132 html_line_breaks(pn_info_string(HEADER_MPN, connector.manufacturer, connector.mpn)),133 html_line_breaks(pn_info_string(HEADER_SPN, connector.supplier, connector.spn))],134 [html_line_breaks(connector.type),135 html_line_breaks(connector.subtype),136 f'{connector.pincount}-pin' if connector.show_pincount else None,137 translate_color(connector.color, self.options.color_mode) if connector.color else None,138 html_colorbar(connector.color)],139 '<!-- connector table -->' if connector.style != 'simple' else None,140 [html_image(connector.image)],141 [html_caption(connector.image)]]142 rows.extend(get_additional_component_table(self, connector))143 rows.append([html_line_breaks(connector.notes)])144 html.extend(nested_html_table(rows, html_bgcolor_attr(connector.bgcolor)))145146 if connector.style != 'simple':147 pinhtml = []148 pinhtml.append('<table border="0" cellspacing="0" cellpadding="3" cellborder="1">')149150 for pinindex, (pinname, pinlabel, pincolor) in enumerate(zip_longest(connector.pins, connector.pinlabels, connector.pincolors)):151 if connector.hide_disconnected_pins and not connector.visible_pins.get(pinname, False):152 continue153 pinhtml.append(' <tr>')154 if connector.ports_left:155 pinhtml.append(f' <td port="p{pinindex+1}l">{pinname}</td>')156 if pinlabel:157 pinhtml.append(f' <td>{pinlabel}</td>')158 if connector.pincolors:159 if pincolor in wv_colors._color_hex.keys():160 pinhtml.append(f' <td sides="tbl">{translate_color(pincolor, self.options.color_mode)}</td>')161 pinhtml.append( ' <td sides="tbr">')162 pinhtml.append( ' <table border="0" cellborder="1"><tr>')163 pinhtml.append(f' <td bgcolor="{wv_colors.translate_color(pincolor, "HEX")}" width="8" height="8" fixedsize="true"></td>')164 pinhtml.append( ' </tr></table>')165 pinhtml.append( ' </td>')166 else:167 pinhtml.append( ' <td colspan="2"></td>')168169 if connector.ports_right:170 pinhtml.append(f' <td port="p{pinindex+1}r">{pinname}</td>')171 pinhtml.append(' </tr>')172173 pinhtml.append(' </table>')174175 html = [row.replace('<!-- connector table -->', '\n'.join(pinhtml)) for row in html]176177 html = '\n'.join(html)178 dot.node(connector.name, label=f'<\n{html}\n>', shape='box', style='filled',179 fillcolor=translate_color(self.options.bgcolor_connector, "HEX"))180181 if len(connector.loops) > 0:182 dot.attr('edge', color='#000000:#ffffff:#000000')183 if connector.ports_left:184 loop_side = 'l'185 loop_dir = 'w'186 elif connector.ports_right:187 loop_side = 'r'188 loop_dir = 'e'189 else:190 raise Exception('No side for loops')191 for loop in connector.loops:192 dot.edge(f'{connector.name}:p{loop[0]}{loop_side}:{loop_dir}',193 f'{connector.name}:p{loop[1]}{loop_side}:{loop_dir}')194195196 # determine if there are double- or triple-colored wires in the harness;197 # if so, pad single-color wires to make all wires of equal thickness198 pad = any(len(colorstr) > 2 for cable in self.cables.values() for colorstr in cable.colors)199200 for cable in self.cables.values():201202 html = []203204 awg_fmt = ''205 if cable.show_equiv:206 # Only convert units we actually know about, i.e. currently207 # mm2 and awg --- other units _are_ technically allowed,208 # and passed through as-is.209 if cable.gauge_unit =='mm\u00B2':210 awg_fmt = f' ({awg_equiv(cable.gauge)} AWG)'211 elif cable.gauge_unit.upper() == 'AWG':212 awg_fmt = f' ({mm2_equiv(cable.gauge)} mm\u00B2)'213214 rows = [[f'{html_bgcolor(cable.bgcolor_title)}{remove_links(cable.name)}'215 if cable.show_name else None],216 [pn_info_string(HEADER_PN, None,217 remove_links(cable.pn)) if not isinstance(cable.pn, list) else None,218 html_line_breaks(pn_info_string(HEADER_MPN,219 cable.manufacturer if not isinstance(cable.manufacturer, list) else None,220 cable.mpn if not isinstance(cable.mpn, list) else None)),221 html_line_breaks(pn_info_string(HEADER_SPN,222 cable.supplier if not isinstance(cable.supplier, list) else None,223 cable.spn if not isinstance(cable.spn, list) else None))],224 [html_line_breaks(cable.type),225 f'{cable.wirecount}x' if cable.show_wirecount else None,226 f'{cable.gauge} {cable.gauge_unit}{awg_fmt}' if cable.gauge else None,227 '+ S' if cable.shield else None,228 f'{cable.length} {cable.length_unit}' if cable.length > 0 else None,229 translate_color(cable.color, self.options.color_mode) if cable.color else None,230 html_colorbar(cable.color)],231 '<!-- wire table -->',232 [html_image(cable.image)],233 [html_caption(cable.image)]]234235 rows.extend(get_additional_component_table(self, cable))236 rows.append([html_line_breaks(cable.notes)])237 html.extend(nested_html_table(rows, html_bgcolor_attr(cable.bgcolor)))238239 wirehtml = []240 wirehtml.append('<table border="0" cellspacing="0" cellborder="0">') # conductor table241 wirehtml.append(' <tr><td> </td></tr>')242243 for i, (connection_color, wirelabel) in enumerate(zip_longest(cable.colors, cable.wirelabels), 1):244 wirehtml.append(' <tr>')245 wirehtml.append(f' <td><!-- {i}_in --></td>')246 wirehtml.append(f' <td>')247248 wireinfo = []249 if cable.show_wirenumbers:250 wireinfo.append(str(i))251 colorstr = wv_colors.translate_color(connection_color, self.options.color_mode)252 if colorstr:253 wireinfo.append(colorstr)254 if cable.wirelabels:255 wireinfo.append(wirelabel if wirelabel is not None else '')256 wirehtml.append(f' {":".join(wireinfo)}')257258 wirehtml.append(f' </td>')259 wirehtml.append(f' <td><!-- {i}_out --></td>')260 wirehtml.append(' </tr>')261262 bgcolors = ['#000000'] + get_color_hex(connection_color, pad=pad) + ['#000000']263 wirehtml.append(f' <tr>')264 wirehtml.append(f' <td colspan="3" border="0" cellspacing="0" cellpadding="0" port="w{i}" height="{(2 * len(bgcolors))}">')265 wirehtml.append(' <table cellspacing="0" cellborder="0" border="0">')266 for j, bgcolor in enumerate(bgcolors[::-1]): # Reverse to match the curved wires when more than 2 colors267 wirehtml.append(f' <tr><td colspan="3" cellpadding="0" height="2" bgcolor="{bgcolor if bgcolor != "" else wv_colors.default_color}" border="0"></td></tr>')268 wirehtml.append(' </table>')269 wirehtml.append(' </td>')270 wirehtml.append(' </tr>')271 if cable.category == 'bundle': # for bundles individual wires can have part information272 # create a list of wire parameters273 wireidentification = []274 if isinstance(cable.pn, list):275 wireidentification.append(pn_info_string(HEADER_PN, None, remove_links(cable.pn[i - 1])))276 manufacturer_info = pn_info_string(HEADER_MPN,277 cable.manufacturer[i - 1] if isinstance(cable.manufacturer, list) else None,278 cable.mpn[i - 1] if isinstance(cable.mpn, list) else None)279 supplier_info = pn_info_string(HEADER_SPN,280 cable.supplier[i - 1] if isinstance(cable.supplier, list) else None,281 cable.spn[i - 1] if isinstance(cable.spn, list) else None)282 if manufacturer_info:283 wireidentification.append(html_line_breaks(manufacturer_info))284 if supplier_info:285 wireidentification.append(html_line_breaks(supplier_info))286 # print parameters into a table row under the wire287 if len(wireidentification) > 0 :288 wirehtml.append(' <tr><td colspan="3">')289 wirehtml.append(' <table border="0" cellspacing="0" cellborder="0"><tr>')290 for attrib in wireidentification:291 wirehtml.append(f' <td>{attrib}</td>')292 wirehtml.append(' </tr></table>')293 wirehtml.append(' </td></tr>')294295 if cable.shield:296 wirehtml.append(' <tr><td> </td></tr>') # spacer297 wirehtml.append(' <tr>')298 wirehtml.append(' <td><!-- s_in --></td>')299 wirehtml.append(' <td>Shield</td>')300 wirehtml.append(' <td><!-- s_out --></td>')301 wirehtml.append(' </tr>')302 if isinstance(cable.shield, str):303 # shield is shown with specified color and black borders304 shield_color_hex = wv_colors.get_color_hex(cable.shield)[0]305 attributes = f'height="6" bgcolor="{shield_color_hex}" border="2" sides="tb"'306 else:307 # shield is shown as a thin black wire308 attributes = f'height="2" bgcolor="#000000" border="0"'309 wirehtml.append(f' <tr><td colspan="3" cellpadding="0" {attributes} port="ws"></td></tr>')310311 wirehtml.append(' <tr><td> </td></tr>')312 wirehtml.append(' </table>')313314 html = [row.replace('<!-- wire table -->', '\n'.join(wirehtml)) for row in html]315316 # connections317 for connection in cable.connections:318 if isinstance(connection.via_port, int): # check if it's an actual wire and not a shield319 dot.attr('edge', color=':'.join(['#000000'] + wv_colors.get_color_hex(cable.colors[connection.via_port - 1], pad=pad) + ['#000000']))320 else: # it's a shield connection321 # shield is shown with specified color and black borders, or as a thin black wire otherwise322 dot.attr('edge', color=':'.join(['#000000', shield_color_hex, '#000000']) if isinstance(cable.shield, str) else '#000000')323 if connection.from_port is not None: # connect to left324 from_connector = self.connectors[connection.from_name]325 from_port = f':p{connection.from_port+1}r' if from_connector.style != 'simple' else ''326 code_left_1 = f'{connection.from_name}{from_port}:e'327 code_left_2 = f'{cable.name}:w{connection.via_port}:w'328 dot.edge(code_left_1, code_left_2)329 if from_connector.show_name:330 from_info = [str(connection.from_name), str(self.connectors[connection.from_name].pins[connection.from_port])]331 if from_connector.pinlabels:332 pinlabel = from_connector.pinlabels[connection.from_port]333 if pinlabel != '':334 from_info.append(pinlabel)335 from_string = ':'.join(from_info)336 else:337 from_string = ''338 html = [row.replace(f'<!-- {connection.via_port}_in -->', from_string) for row in html]339 if connection.to_port is not None: # connect to right340 to_connector = self.connectors[connection.to_name]341 code_right_1 = f'{cable.name}:w{connection.via_port}:e'342 to_port = f':p{connection.to_port+1}l' if self.connectors[connection.to_name].style != 'simple' else ''343 code_right_2 = f'{connection.to_name}{to_port}:w'344 dot.edge(code_right_1, code_right_2)345 if to_connector.show_name:346 to_info = [str(connection.to_name), str(self.connectors[connection.to_name].pins[connection.to_port])]347 if to_connector.pinlabels:348 pinlabel = to_connector.pinlabels[connection.to_port]349 if pinlabel != '':350 to_info.append(pinlabel)351 to_string = ':'.join(to_info)352 else:353 to_string = ''354 html = [row.replace(f'<!-- {connection.via_port}_out -->', to_string) for row in html]355356 style, bgcolor = ('filled,dashed', self.options.bgcolor_bundle) if cable.category == 'bundle' else \357 ('filled', self.options.bgcolor_cable)358 html = '\n'.join(html)359 dot.node(cable.name, label=f'<\n{html}\n>', shape='box',360 style=style, fillcolor=translate_color(bgcolor, "HEX"))361362 def typecheck(name: str, value: Any, expect: type) -> None:363 if not isinstance(value, expect):364 raise Exception(f'Unexpected value type of {name}: Expected {expect}, got {type(value)}\n{value}')365366 # TODO?: Differ between override attributes and HTML?367 if self.tweak.override is not None:368 typecheck('tweak.override', self.tweak.override, dict)369 for k, d in self.tweak.override.items():370 typecheck(f'tweak.override.{k} key', k, str)371 typecheck(f'tweak.override.{k} value', d, dict)372 for a, v in d.items():373 typecheck(f'tweak.override.{k}.{a} key', a, str)374 typecheck(f'tweak.override.{k}.{a} value', v, (str, type(None)))375376 # Override generated attributes of selected entries matching tweak.override.377 for i, entry in enumerate(dot.body):378 if isinstance(entry, str):379 # Find a possibly quoted keyword after leading TAB(s) and followed by [ ].380 match = re.match(r'^\t*(")?((?(1)[^"]|[^ "])+)(?(1)") \[.*\]$', entry, re.S)381 keyword = match and match[2]382 if keyword in self.tweak.override.keys():383 for attr, value in self.tweak.override[keyword].items():384 if value is None:385 entry, n_subs = re.subn(f'( +)?{attr}=("[^"]*"|[^] ]*)(?(1)| *)', '', entry)386 if n_subs < 1:387 sys.stderr.write(f'Harness.create_graph() warning: {attr} not found in {keyword}!')388 elif n_subs > 1:389 sys.stderr.write(f'Harness.create_graph() warning: {attr} removed {n_subs} times in {keyword}!')390 continue391392 if len(value) == 0 or ' ' in value:393 value = value.replace('"', r'\"')394 value = f'"{value}"'395 entry, n_subs = re.subn(f'{attr}=("[^"]*"|[^] ]*)', f'{attr}={value}', entry)396 if n_subs < 1:397 # If attr not found, then append it398 entry = re.sub(r'\]$', f' {attr}={value}]', entry)399 elif n_subs > 1:400 sys.stderr.write(f'Harness.create_graph() warning: {attr} overridden {n_subs} times in {keyword}!')401402 dot.body[i] = entry403404 if self.tweak.append is not None:405 if isinstance(self.tweak.append, list):406 for i, element in enumerate(self.tweak.append, 1):407 typecheck(f'tweak.append[{i}]', element, str)408 dot.body.extend(self.tweak.append)409 else:410 typecheck('tweak.append', self.tweak.append, str)411 dot.body.append(self.tweak.append)412413 return dot414415416 def png(self):417 from io import BytesIO418 graph = self.create_graph()419 data = BytesIO()420 data.write(graph.pipe(format='png'))421 data.seek(0)422 return data.read()423424425 def svg(self):426 from io import BytesIO427 graph = self.create_graph()428 data = BytesIO()429 data.write(graph.pipe(format='svg'))430 data.seek(0)431 return data.read()432433 def output(self, filename: (str, Path), view: bool = False, cleanup: bool = True, fmt: tuple = ('pdf', )) -> None:434 # graphical output435 graph = self.create_graph()436 for f in fmt:437 graph.format = f438 graph.render(filename=filename, view=view, cleanup=cleanup)439440 graph.save(filename=f'{filename}.gv')441 # bom output442 bomlist = bom_list(self.bom())443 with open_file_write(f'{filename}.bom.tsv') as file:444 file.write(tuplelist2tsv(bomlist))445 # HTML output446 generate_html_output(filename, bomlist, self.metadata, self.options)447448 def bom(self):449 if not self._bom:450 self._bom = generate_bom(self)451 return self._bom