328 lines
12 KiB
Python
328 lines
12 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import Any, List
|
|
from graphviz import Graph
|
|
|
|
COLOR_CODES = {'DIN': ['WH','BN','GN','YE','GY','PK','BU','RD','BK','VT'], # ,'GYPK','RDBU','WHGN','BNGN','WHYE','YEBN','WHGY','GYBN','WHPK','PKBN'],
|
|
'IEC': ['BN','RD','OG','YE','GN','BU','VT','GY','WH','BK'],
|
|
'BW': ['BK','WH']}
|
|
|
|
# TODO: parse and render double-colored cables ('RDBU' etc)
|
|
color_hex = {
|
|
'BK': '#000000',
|
|
'WH': '#ffffff',
|
|
'GY': '#999999',
|
|
'PK': '#ff66cc',
|
|
'RD': '#ff0000',
|
|
'OG': '#ff8000',
|
|
'YE': '#ffff00',
|
|
'GN': '#009900',
|
|
'TQ': '#00ffff',
|
|
'BU': '#0066ff',
|
|
'VT': '#8000ff',
|
|
'BN': '#666600',
|
|
}
|
|
|
|
color_full = {
|
|
'BK': 'black',
|
|
'WH': 'white',
|
|
'GY': 'grey',
|
|
'PK': 'pink',
|
|
'RD': 'red',
|
|
'OG': 'orange',
|
|
'YE': 'yellow',
|
|
'GN': 'green',
|
|
'TQ': 'turquoise',
|
|
'BU': 'blue',
|
|
'VT': 'violet',
|
|
'BN': 'brown',
|
|
}
|
|
|
|
color_ger = {
|
|
'BK': 'sw',
|
|
'WH': 'ws',
|
|
'GY': 'gr',
|
|
'PK': 'rs',
|
|
'RD': 'rt',
|
|
'OG': 'or',
|
|
'YE': 'ge',
|
|
'GN': 'gn',
|
|
'TQ': 'tk',
|
|
'BU': 'bl',
|
|
'VT': 'vi',
|
|
'BN': 'br',
|
|
}
|
|
|
|
class Harness:
|
|
|
|
def __init__(self):
|
|
self.color_mode = 'SHORT'
|
|
self.nodes = {}
|
|
self.cables = {}
|
|
|
|
def add_node(self, name, *args, **kwargs):
|
|
self.nodes[name] = Node(name, *args, **kwargs)
|
|
|
|
def add_cable(self, name, *args, **kwargs):
|
|
self.cables[name] = Cable(name, *args, **kwargs)
|
|
|
|
def loop(self, node_name, from_pin, to_pin):
|
|
self.nodes[node_name].loop(from_pin, to_pin)
|
|
|
|
def connect(self, from_name, from_pin, via_name, via_pin, to_name, to_pin):
|
|
self.cables[via_name].connect(from_name, from_pin, via_pin, to_name, to_pin)
|
|
|
|
def connect_all_straight(self, cable_name, from_name, to_name):
|
|
self.cables[cable_name].connect_all_straight(from_name, to_name)
|
|
|
|
def create_graph(self):
|
|
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='transparent', fontname=font)
|
|
dot.attr('node', shape='record', style='rounded,filled', fillcolor='white', fontname=font)
|
|
dot.attr('edge', style='bold', fontname=font)
|
|
|
|
# prepare ports on connectors depending on which side they will connect
|
|
for k, c in self.cables.items():
|
|
for x in c.connections:
|
|
if x[1] is not None: # connect to left
|
|
self.nodes[x[0]].ports_right = True
|
|
if x[4] is not None: # connect to right
|
|
self.nodes[x[3]].ports_left = True
|
|
|
|
for k, n in self.nodes.items():
|
|
# a = attributes
|
|
a = [n.type,
|
|
n.gender,
|
|
'{}-pin'.format(len(n.pinout)) if n.show_num_pins else '']
|
|
# p = pinout
|
|
p = [[],[],[]]
|
|
p[1] = list(n.pinout)
|
|
for i, x in enumerate(n.pinout, 1):
|
|
if n.ports_left:
|
|
p[0].append('<p{portno}l>{portno}'.format(portno=i))
|
|
if n.ports_right:
|
|
p[2].append('<p{portno}r>{portno}'.format(portno=i))
|
|
# l = label
|
|
l = [n.name if n.show_name else '', a, p]
|
|
dot.node(k, label=nested(l))
|
|
|
|
if len(n.loops) > 0:
|
|
dot.attr('edge',color='#000000')
|
|
if n.ports_left:
|
|
loop_side = 'l'
|
|
loop_dir = 'w'
|
|
elif n.ports_right:
|
|
loop_side = 'r'
|
|
loop_dir = 'e'
|
|
else:
|
|
raise Exception('No side for loops')
|
|
for x in n.loops:
|
|
dot.edge('{name}:p{port_from}{loop_side}:{loop_dir}'.format(name=n.name, port_from=x[0], port_to=x[1], loop_side=loop_side, loop_dir=loop_dir),
|
|
'{name}:p{port_to}{loop_side}:{loop_dir}'.format(name=n.name, port_from=x[0], port_to=x[1], loop_side=loop_side, loop_dir=loop_dir))
|
|
|
|
for k, c in self.cables.items():
|
|
# a = attributes
|
|
a = ['{}x'.format(len(c.colors)) if c.show_num_wires else '',
|
|
'{} mm\u00B2{}'.format(c.mm2, ' ({} AWG)'.format(awg_equiv(c.mm2)) if c.show_equiv else '') if c.mm2 is not None else '',
|
|
c.awg,
|
|
'+ S' if c.shield else '',
|
|
'{} m'.format(c.length) if c.length > 0 else '']
|
|
# p = pinout
|
|
p = [[],[],[]]
|
|
for i, x in enumerate(c.colors,1):
|
|
if c.show_pinout:
|
|
p[0].append('<w{wireno}i>{wireno}'.format(wireno=i))
|
|
p[1].append('{wirecolor}'.format(wirecolor=translate_color(x, self.color_mode)))
|
|
p[2].append('<w{wireno}o>{wireno}'.format(wireno=i))
|
|
else:
|
|
p[1].append('<w{wireno}>{wirecolor}'.format(wireno=i,wirecolor=translate_color(x, self.color_mode)))
|
|
if c.shield:
|
|
if c.show_pinout:
|
|
p[0].append('<wsi>')
|
|
p[1].append('Shield')
|
|
p[2].append('<wso>')
|
|
else:
|
|
p[1].append('<ws>Shield')
|
|
# l = label
|
|
l = [c.name if c.show_name else '', a, p]
|
|
dot.node(k, label=nested(l))
|
|
|
|
# connections
|
|
for x in c.connections:
|
|
if isinstance(x[2], int): # check if it's an actual wire and not a shield
|
|
search_color = c.colors[x[2]-1]
|
|
if search_color in color_hex:
|
|
dot.attr('edge',color='#000000:{wire_color}:#000000'.format(wire_color=color_hex[search_color]))
|
|
else: # color name not found
|
|
dot.attr('edge',color='#000000')
|
|
else: # it's a shield connection
|
|
dot.attr('edge',color='#000000')
|
|
if x[1] is not None: # connect to left
|
|
dot.edge('{from_name}:p{from_port}r'.format(from_name=x[0],from_port=x[1]),
|
|
'{via_name}:w{via_wire}{via_subport}'.format(via_name=c.name, via_wire=x[2], via_subport='i' if c.show_pinout else ''))
|
|
# self.nodes[x[0]].ports_right = True
|
|
if x[4] is not None: # connect to right
|
|
dot.edge('{via_name}:w{via_wire}{via_subport}'.format(via_name=c.name, via_wire=x[2], via_subport='o' if c.show_pinout else ''),
|
|
'{to_name}:p{to_port}l'.format(to_name=x[3], to_port=x[4]))
|
|
# self.nodes[x[3]].ports_left = True
|
|
|
|
return dot
|
|
|
|
def output(self, filename, directory='_output', view=False, cleanup=True, format='pdf'):
|
|
d = self.create_graph()
|
|
for f in format:
|
|
d.format = f
|
|
d.render(filename=filename, directory=directory, view=view, cleanup=cleanup)
|
|
d.save(filename='{}.gv'.format(filename), directory=directory)
|
|
|
|
@dataclass
|
|
class Node:
|
|
name: str
|
|
type: str = None
|
|
gender: str = None
|
|
num_pins: int = None
|
|
pinout: List[Any] = field(default_factory=list)
|
|
show_name: bool = False
|
|
show_num_pins: bool = False
|
|
|
|
def __post_init__(self):
|
|
self.ports_left = False
|
|
self.ports_right = False
|
|
self.loops = []
|
|
|
|
if self.pinout:
|
|
if self.num_pins is not None:
|
|
raise Exception('You cannot specify both pinout and num_pins')
|
|
else:
|
|
if not self.num_pins:
|
|
self.num_pins = 1
|
|
self.pinout = ['',] * self.num_pins
|
|
|
|
def loop(self, from_pin, to_pin):
|
|
self.loops.append((from_pin, to_pin))
|
|
|
|
@dataclass
|
|
class Cable:
|
|
name: str
|
|
mm2: float = None
|
|
awg: int = None
|
|
show_equiv: bool = False
|
|
length: float = 0
|
|
num_wires: int = None
|
|
shield: bool = False
|
|
colors: List[Any] = field(default_factory=list)
|
|
color_code: str = None
|
|
show_name: bool = False
|
|
show_pinout: bool = False
|
|
show_num_wires: bool = True
|
|
|
|
def __post_init__(self):
|
|
if self.mm2 and self.awg:
|
|
raise Exception('You cannot define both mm2 and awg!')
|
|
self.connections = []
|
|
|
|
if self.num_wires: # number of wires explicitly defined
|
|
if self.colors: # use custom color palette (partly or looped if needed)
|
|
pass
|
|
elif self.color_code: # use standard color palette (partly or looped if needed)
|
|
if self.color_code not in COLOR_CODES:
|
|
raise Exception('Unknown color code')
|
|
self.colors = COLOR_CODES[self.color_code]
|
|
else: # no colors defined, add dummy colors
|
|
self.colors = [''] * self.num_wires
|
|
|
|
# make color code loop around if more wires than colors
|
|
if self.num_wires > len(self.colors):
|
|
m = self.num_wires // len(self.colors) + 1
|
|
self.colors = self.colors * int(m)
|
|
# cut off excess after looping
|
|
self.colors = self.colors[:self.num_wires]
|
|
|
|
else: # num_wires implicit in length of color list
|
|
if not self.colors:
|
|
raise Exception('Unknown number of wires. Must specify num_wires or colors (implicit length)')
|
|
self.num_wires = len(self.colors)
|
|
|
|
def connect(self, from_name, from_pin, via_pin, to_name, to_pin):
|
|
from_pin = int2tuple(from_pin)
|
|
via_pin = int2tuple(via_pin)
|
|
to_pin = int2tuple(to_pin)
|
|
if len(from_pin) != len(to_pin):
|
|
raise Exception('from_pin must have the same number of elements as to_pin')
|
|
for i, x in enumerate(from_pin):
|
|
self.connections.append((from_name, from_pin[i], via_pin[i], to_name, to_pin[i]))
|
|
|
|
def connect_all_straight(self, from_name, to_name):
|
|
self.connect(from_name, 'auto', 'auto', to_name, 'auto')
|
|
|
|
def nested(input):
|
|
l = []
|
|
for x in input:
|
|
if isinstance(x, list):
|
|
if len(x) > 0:
|
|
n = nested(x)
|
|
if n != '':
|
|
l.append('{' + n + '}')
|
|
else:
|
|
if x is not None:
|
|
if x != '':
|
|
l.append(str(x))
|
|
s = '|'.join(l)
|
|
return s
|
|
|
|
def int2tuple(input):
|
|
if isinstance(input, tuple):
|
|
output = input
|
|
else:
|
|
output = (input,)
|
|
return output
|
|
|
|
def translate_color(input, color_mode):
|
|
if input == '':
|
|
output = ''
|
|
else:
|
|
if color_mode == 'full':
|
|
output = color_full[input].lower()
|
|
elif color_mode == 'FULL':
|
|
output = color_hex[input].upper()
|
|
elif color_mode == 'hex':
|
|
output = color_hex[input].lower()
|
|
elif color_mode == 'HEX':
|
|
output = color_hex[input].upper()
|
|
elif color_mode == 'ger':
|
|
output = color_ger[input].lower()
|
|
elif color_mode == 'GER':
|
|
output = color_ger[input].upper()
|
|
elif color_mode == 'short':
|
|
output = input.lower()
|
|
elif color_mode == 'SHORT':
|
|
output = input.upper()
|
|
else:
|
|
raise Exception('Unknown color mode')
|
|
return output
|
|
|
|
def awg_equiv(mm2):
|
|
awg_equiv_table = {
|
|
'0.09': 28,
|
|
'0.14': 26,
|
|
'0.25': 24,
|
|
'0.34': 22,
|
|
'0.5': 21,
|
|
'0.75': 20,
|
|
'1': 18,
|
|
'1.5': 16,
|
|
'2.5': 14,
|
|
'4': 12,
|
|
'6': 10,
|
|
'10': 8,
|
|
'16': 6,
|
|
'25': 4,
|
|
}
|
|
k = str(mm2)
|
|
if k in awg_equiv_table:
|
|
return awg_equiv_table[k]
|
|
else:
|
|
return None
|