#!/usr/bin/env python # -*- coding: utf-8 -*- from wireviz.DataClasses import Connector, Cable from graphviz import Graph from wireviz import wv_colors, wv_helper from wireviz.wv_colors import get_color_hex from wireviz.wv_helper import awg_equiv, mm2_equiv, tuplelist2tsv, \ nested_html_table, flatten2d, index_if_list, html_line_breaks, \ graphviz_line_breaks, remove_line_breaks, open_file_read, open_file_write, \ manufacturer_info_field from collections import Counter from typing import List from pathlib import Path import re class Harness: def __init__(self): self.color_mode = 'SHORT' self.connectors = {} self.cables = {} self.additional_bom_items = [] def add_connector(self, name: str, *args, **kwargs) -> None: self.connectors[name] = Connector(name, *args, **kwargs) def add_cable(self, name: str, *args, **kwargs) -> None: self.cables[name] = Cable(name, *args, **kwargs) def add_bom_item(self, item: dict) -> None: self.additional_bom_items.append(item) def connect(self, from_name: str, from_pin: (int, str), via_name: str, via_pin: (int, str), to_name: str, to_pin: (int, str)) -> None: for (name, pin) in zip([from_name, to_name], [from_pin, to_pin]): # check from and to connectors if name is not None and name in self.connectors: connector = self.connectors[name] if pin in connector.pins and pin in connector.pinlabels: if connector.pins.index(pin) == connector.pinlabels.index(pin): # TODO: Maybe issue a warning? It's not worthy of an exception if it's unambiguous, but maybe risky? pass else: raise Exception(f'{name}:{pin} is defined both in pinlabels and pins, for different pins.') if pin in connector.pinlabels: if connector.pinlabels.count(pin) > 1: raise Exception(f'{name}:{pin} is defined more than once.') else: index = connector.pinlabels.index(pin) pin = connector.pins[index] # map pin name to pin number if name == from_name: from_pin = pin if name == to_name: to_pin = pin if not pin in connector.pins: raise Exception(f'{name}:{pin} not found.') self.cables[via_name].connect(from_name, from_pin, via_pin, to_name, to_pin) if from_name in self.connectors: self.connectors[from_name].activate_pin(from_pin) if to_name in self.connectors: self.connectors[to_name].activate_pin(to_pin) def create_graph(self) -> Graph: dot = Graph() dot.body.append('// Graph generated by WireViz') dot.body.append('// https://github.com/formatc1702/WireViz') font = 'arial' dot.attr('graph', rankdir='LR', ranksep='2', bgcolor='white', nodesep='0.33', fontname=font) dot.attr('node', shape='record', style='filled', fillcolor='white', fontname=font) dot.attr('edge', style='bold', fontname=font) # prepare ports on connectors depending on which side they will connect for _, cable in self.cables.items(): for connection_color in cable.connections: if connection_color.from_port is not None: # connect to left self.connectors[connection_color.from_name].ports_right = True if connection_color.to_port is not None: # connect to right self.connectors[connection_color.to_name].ports_left = True for key, connector in self.connectors.items(): rows = [[connector.name if connector.show_name else None], [f'P/N: {connector.pn}' if connector.pn else None, manufacturer_info_field(connector.manufacturer, connector.mpn)], [html_line_breaks(connector.type), html_line_breaks(connector.subtype), f'{connector.pincount}-pin' if connector.show_pincount else None, connector.color, '' if connector.color else None], '' if connector.style != 'simple' else None, [html_line_breaks(connector.notes)]] html = nested_html_table(rows) if connector.color: # add color bar next to color info, if present colorbar = f' bgcolor="{wv_colors.translate_color(connector.color, "HEX")}" width="4">' # leave out ' tag html = html.replace('>', colorbar) if connector.style != 'simple': pinlist = [] for pin, pinlabel in zip(connector.pins, connector.pinlabels): if connector.hide_disconnected_pins and not connector.visible_pins.get(pin, False): continue pinlist.append([f'{pin}' if connector.ports_left else None, f'{pinlabel}' if pinlabel else '', f'{pin}' if connector.ports_right else None]) pinhtml = '' for i, pin in enumerate(pinlist): pinhtml = f'{pinhtml}' for column in pin: if column is not None: pinhtml = f'{pinhtml}{column}' pinhtml = f'{pinhtml}' pinhtml = f'{pinhtml}
' html = html.replace('', pinhtml) dot.node(key, label=f'<{html}>', shape='none', margin='0', style='filled', fillcolor='white') if len(connector.loops) > 0: dot.attr('edge', color='#000000:#ffffff:#000000') if connector.ports_left: loop_side = 'l' loop_dir = 'w' elif connector.ports_right: loop_side = 'r' loop_dir = 'e' else: raise Exception('No side for loops') for loop in connector.loops: dot.edge(f'{connector.name}:p{loop[0]}{loop_side}:{loop_dir}', f'{connector.name}:p{loop[1]}{loop_side}:{loop_dir}') for _, cable in self.cables.items(): awg_fmt = '' if cable.show_equiv: # Only convert units we actually know about, i.e. currently # mm2 and awg --- other units _are_ technically allowed, # and passed through as-is. if cable.gauge_unit =='mm\u00B2': awg_fmt = f' ({awg_equiv(cable.gauge)} AWG)' elif cable.gauge_unit.upper() == 'AWG': awg_fmt = f' ({mm2_equiv(cable.gauge)} mm\u00B2)' identification = [f'P/N: {cable.pn}' if (cable.pn and not isinstance(cable.pn, list)) else '', manufacturer_info_field(cable.manufacturer if not isinstance(cable.manufacturer, list) else None, cable.mpn if not isinstance(cable.mpn, list) else None)] identification = list(filter(None, identification)) attributes = [html_line_breaks(cable.type) if cable.type else '', f'{len(cable.colors)}x' if cable.show_wirecount else '', f'{cable.gauge} {cable.gauge_unit}{awg_fmt}' if cable.gauge else '', '+ S' if cable.shield else '', f'{cable.length} m' if cable.length > 0 else ''] attributes = list(filter(None, attributes)) html = '' # main table if cable.show_name or len(attributes) > 0: html = f'{html}' # name+attributes table html = f'{html}' # spacer between attributes and wires html = f'{html}' # main table if cable.notes: html = f'{html}' # notes table html = f'{html}' # spacer at the end html = f'{html}
' # name+attributes table if cable.show_name: html = f'{html}' if(len(identification) > 0): # print an identification row if values specified html = f'{html}' # end identification row if(len(attributes) > 0): html = f'{html}' # attribute row for attrib in attributes: html = f'{html}' html = f'{html}' # attribute row html = f'{html}
{cable.name}
' for attrib in identification[0:-1]: html = f'{html}' # all columns except last have a border on the right (sides="R") if len(identification) > 0: html = f'{html}' # last column has no border on the right because the enclosing table borders it html = f'{html}
{attrib}{identification[-1]}
{attrib}
 
' # conductor table # determine if there are double- or triple-colored wires; # if so, pad single-color wires to make all wires of equal thickness colorlengths = list(map(len, cable.colors)) pad = 4 in colorlengths or 6 in colorlengths for i, connection_color in enumerate(cable.colors, 1): p = [] p.append(f'') p.append(wv_colors.translate_color(connection_color, self.color_mode)) p.append(f'') html = f'{html}' for bla in p: html = f'{html}' html = f'{html}' bgcolors = ['#000000'] + get_color_hex(connection_color, pad=pad) + ['#000000'] html = f'{html}' if(cable.category == 'bundle'): # for bundles individual wires can have part information # create a list of wire parameters wireidentification = [] if isinstance(cable.pn, list): wireidentification.append(f'P/N: {cable.pn[i - 1]}') manufacturer_info = manufacturer_info_field(cable.manufacturer[i - 1] if isinstance(cable.manufacturer, list) else None, cable.mpn[i - 1] if isinstance(cable.mpn, list) else None) if manufacturer_info: wireidentification.append(manufacturer_info) # print parameters into a table row under the wire if(len(wireidentification) > 0): html = f'{html}' if cable.shield: p = ['', 'Shield', ''] html = f'{html}' # spacer html = f'{html}' for bla in p: html = html + f'' html = f'{html}' if isinstance(cable.shield, str): # shield is shown with specified color and black borders shield_color_hex = wv_colors.get_color_hex(cable.shield)[0] attributes = f'height="6" bgcolor="{shield_color_hex}" border="2" sides="tb"' else: # shield is shown as a thin black wire attributes = f'height="2" bgcolor="#000000" border="0"' html = f'{html}' html = f'{html}' # spacer at the end html = f'{html}
{bla}
' for j, bgcolor in enumerate(bgcolors[::-1]): # Reverse to match the curved wires when more than 2 colors html = f'{html}' html = html + '
' for attrib in wireidentification: html = f'{html}' html = f'{html}
{attrib}
 
{bla}
 
' # conductor table html = f'{html}
{html_line_breaks(cable.notes)}
 
' # main table # connections for connection_color in cable.connections: if isinstance(connection_color.via_port, int): # check if it's an actual wire and not a shield dot.attr('edge', color=':'.join(['#000000'] + wv_colors.get_color_hex(cable.colors[connection_color.via_port - 1], pad=pad) + ['#000000'])) else: # it's a shield connection # shield is shown with specified color and black borders, or as a thin black wire otherwise dot.attr('edge', color=':'.join(['#000000', shield_color_hex, '#000000']) if isinstance(cable.shield, str) else '#000000') if connection_color.from_port is not None: # connect to left from_port = f':p{connection_color.from_port}r' if self.connectors[connection_color.from_name].style != 'simple' else '' code_left_1 = f'{connection_color.from_name}{from_port}:e' code_left_2 = f'{cable.name}:w{connection_color.via_port}:w' dot.edge(code_left_1, code_left_2) from_string = f'{connection_color.from_name}:{connection_color.from_port}' if self.connectors[connection_color.from_name].show_name else '' html = html.replace(f'', from_string) if connection_color.to_port is not None: # connect to right code_right_1 = f'{cable.name}:w{connection_color.via_port}:e' to_port = f':p{connection_color.to_port}l' if self.connectors[connection_color.to_name].style != 'simple' else '' code_right_2 = f'{connection_color.to_name}{to_port}:w' dot.edge(code_right_1, code_right_2) to_string = f'{connection_color.to_name}:{connection_color.to_port}' if self.connectors[connection_color.to_name].show_name else '' html = html.replace(f'', to_string) dot.node(cable.name, label=f'<{html}>', shape='box', style='filled,dashed' if cable.category == 'bundle' else '', margin='0', fillcolor='white') return dot @property def png(self): from io import BytesIO graph = self.create_graph() data = BytesIO() data.write(graph.pipe(format='png')) data.seek(0) return data.read() @property def svg(self): from io import BytesIO graph = self.create_graph() data = BytesIO() data.write(graph.pipe(format='svg')) data.seek(0) return data.read() def output(self, filename: (str, Path), view: bool = False, cleanup: bool = True, fmt: tuple = ('pdf', )) -> None: # graphical output graph = self.create_graph() for f in fmt: graph.format = f graph.render(filename=filename, view=view, cleanup=cleanup) graph.save(filename=f'{filename}.gv') # bom output bom_list = self.bom_list() with open_file_write(f'{filename}.bom.tsv') as file: file.write(tuplelist2tsv(bom_list)) # HTML output with open_file_write(f'{filename}.html') as file: file.write('\n') file.write('') file.write('

Diagram

') with open_file_read(f'{filename}.svg') as svg: file.write(re.sub( '^<[?]xml [^?>]*[?]>[^<]*]*>', '', svg.read(1024), 1)) for svgdata in svg: file.write(svgdata) file.write('

Bill of Materials

') listy = flatten2d(bom_list) file.write('') file.write('') for item in listy[0]: file.write(f'') file.write('') for row in listy[1:]: file.write('') for i, item in enumerate(row): item_str = item.replace('\u00b2', '²') align = 'align="right"' if listy[0][i] == 'Qty' else '' file.write(f'') file.write('') file.write('
{item}
{item_str}
') file.write('') def bom(self): bom = [] bom_connectors = [] bom_cables = [] bom_extra = [] # connectors connector_group = lambda c: (c.type, c.subtype, c.pincount, c.manufacturer, c.mpn, c.pn) for group in Counter([connector_group(v) for v in self.connectors.values()]): items = {k: v for k, v in self.connectors.items() if connector_group(v) == group} shared = next(iter(items.values())) designators = list(items.keys()) designators.sort() conn_type = f', {remove_line_breaks(shared.type)}' if shared.type else '' conn_subtype = f', {remove_line_breaks(shared.subtype)}' if shared.subtype else '' conn_pincount = f', {shared.pincount} pins' if shared.style != 'simple' else '' conn_color = f', {shared.color}' if shared.color else '' name = f'Connector{conn_type}{conn_subtype}{conn_pincount}{conn_color}' item = {'item': name, 'qty': len(designators), 'unit': '', 'designators': designators if shared.show_name else '', 'manufacturer': shared.manufacturer, 'mpn': shared.mpn, 'pn': shared.pn} bom_connectors.append(item) bom_connectors = sorted(bom_connectors, key=lambda k: k['item']) # https://stackoverflow.com/a/73050 bom.extend(bom_connectors) # cables # TODO: If category can have other non-empty values than 'bundle', maybe it should be part of item name? # The category needs to be included in cable_group to keep the bundles excluded. cable_group = lambda c: (c.category, c.type, c.gauge, c.gauge_unit, c.wirecount, c.shield, c.manufacturer, c.mpn, c.pn) for group in Counter([cable_group(v) for v in self.cables.values() if v.category != 'bundle']): items = {k: v for k, v in self.cables.items() if cable_group(v) == group} shared = next(iter(items.values())) designators = list(items.keys()) designators.sort() total_length = sum(i.length for i in items.values()) cable_type = f', {remove_line_breaks(shared.type)}' if shared.type else '' gauge_name = f' x {shared.gauge} {shared.gauge_unit}' if shared.gauge else ' wires' shield_name = ' shielded' if shared.shield else '' name = f'Cable{cable_type}, {shared.wirecount}{gauge_name}{shield_name}' item = {'item': name, 'qty': round(total_length, 3), 'unit': 'm', 'designators': designators, 'manufacturer': shared.manufacturer, 'mpn': shared.mpn, 'pn': shared.pn} bom_cables.append(item) # bundles (ignores wirecount) wirelist = [] # list all cables again, since bundles are represented as wires internally, with the category='bundle' set for bundle in self.cables.values(): if bundle.category == 'bundle': # add each wire from each bundle to the wirelist for index, color in enumerate(bundle.colors, 0): wirelist.append({'type': bundle.type, 'gauge': bundle.gauge, 'gauge_unit': bundle.gauge_unit, 'length': bundle.length, 'color': color, 'designator': bundle.name, 'manufacturer': index_if_list(bundle.manufacturer, index), 'mpn': index_if_list(bundle.mpn, index), 'pn': index_if_list(bundle.pn, index)}) # join similar wires from all the bundles to a single BOM item wire_group = lambda w: (w.get('type', None), w['gauge'], w['gauge_unit'], w['color'], w['manufacturer'], w['mpn'], w['pn']) for group in Counter([wire_group(v) for v in wirelist]): items = [v for v in wirelist if wire_group(v) == group] shared = items[0] designators = [i['designator'] for i in items] designators = list(dict.fromkeys(designators)) # remove duplicates designators.sort() total_length = sum(i['length'] for i in items) wire_type = f', {remove_line_breaks(shared["type"])}' if shared.get('type', None) else '' gauge_name = f', {shared["gauge"]} {shared["gauge_unit"]}' if shared.get('gauge', None) else '' gauge_color = f', {shared["color"]}' if 'color' in shared != '' else '' name = f'Wire{wire_type}{gauge_name}{gauge_color}' item = {'item': name, 'qty': round(total_length, 3), 'unit': 'm', 'designators': designators, 'manufacturer': shared['manufacturer'], 'mpn': shared['mpn'], 'pn': shared['pn']} bom_cables.append(item) bom_cables = sorted(bom_cables, key=lambda k: k['item']) # sort list of dicts by their values (https://stackoverflow.com/a/73050) bom.extend(bom_cables) for item in self.additional_bom_items: name = item['description'] if item.get('description', None) else '' if isinstance(item.get('designators', None), List): item['designators'].sort() # sort designators if a list is provided item = {'item': name, 'qty': item.get('qty', None), 'unit': item.get('unit', None), 'designators': item.get('designators', None), 'manufacturer': item.get('manufacturer', None), 'mpn': item.get('mpn', None), 'pn': item.get('pn', None)} bom_extra.append(item) bom_extra = sorted(bom_extra, key=lambda k: k['item']) bom.extend(bom_extra) return bom def bom_list(self): bom = self.bom() keys = ['item', 'qty', 'unit', 'designators'] # these BOM columns will always be included for fieldname in ['pn', 'manufacturer', 'mpn']: # these optional BOM columns will only be included if at least one BOM item actually uses them if any(fieldname in x and x.get(fieldname, None) for x in bom): keys.append(fieldname) bom_list = [] # list of staic bom header names, headers not specified here are generated by capitilising the internal name bom_headings = { "pn": "P/N", "mpn": "MPN" } bom_list.append([(bom_headings[k] if k in bom_headings else k.capitalize()) for k in keys]) # create header row with keys for item in bom: item_list = [item.get(key, '') for key in keys] # fill missing values with blanks item_list = [', '.join(subitem) if isinstance(subitem, List) else subitem for subitem in item_list] # convert any lists into comma separated strings item_list = ['' if subitem is None else subitem for subitem in item_list] # if a field is missing for some (but not all) BOM items bom_list.append(item_list) return bom_list