400 lines
14 KiB
Python
400 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Dict, List
|
|
|
|
from graphviz import Graph
|
|
|
|
import wireviz.wv_colors
|
|
from wireviz.wv_bom import bom_list
|
|
from wireviz.wv_dataclasses import (
|
|
Arrow,
|
|
ArrowWeight,
|
|
BomCategory,
|
|
Cable,
|
|
Component,
|
|
Connector,
|
|
MateComponent,
|
|
MatePin,
|
|
Metadata,
|
|
Options,
|
|
Side,
|
|
)
|
|
from wireviz.wv_graphviz import (
|
|
calculate_node_bgcolor,
|
|
gv_connector_loops,
|
|
gv_edge_mate,
|
|
gv_edge_wire,
|
|
gv_node_component,
|
|
parse_arrow_str,
|
|
set_dot_basics,
|
|
)
|
|
from wireviz.wv_output import embed_svg_images_file, generate_html_output
|
|
from wireviz.wv_utils import bom2tsv
|
|
|
|
|
|
@dataclass
|
|
class Harness:
|
|
metadata: Metadata
|
|
options: Options
|
|
additional_bom_items: List[Component] = field(default_factory=list)
|
|
shared_bom: Dict = field(default_factory=dict)
|
|
|
|
def __post_init__(self):
|
|
self.connectors = {}
|
|
self.cables = {}
|
|
self.mates = []
|
|
self.bom = {}
|
|
self.additional_bom_items = []
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
pn = self.metadata.get("pn", "")
|
|
output_name = self.metadata["output_name"]
|
|
return pn + output_name
|
|
|
|
def add_connector(self, designator: str, *args, **kwargs) -> None:
|
|
conn = Connector(designator=designator, *args, **kwargs)
|
|
self.connectors[designator] = conn
|
|
|
|
def add_cable(self, designator: str, *args, **kwargs) -> None:
|
|
cbl = Cable(designator=designator, *args, **kwargs)
|
|
self.cables[designator] = cbl
|
|
|
|
def add_additional_bom_item(self, item: dict) -> None:
|
|
new_item = Component(**item, category=BomCategory.ADDITIONAL)
|
|
self.additional_bom_items.append(new_item)
|
|
|
|
def add_mate_pin(self, from_name, from_pin, to_name, to_pin, arrow_str) -> None:
|
|
from_con = self.connectors[from_name]
|
|
from_pin_obj = from_con.pin_objects[from_pin]
|
|
to_con = self.connectors[to_name]
|
|
to_pin_obj = to_con.pin_objects[to_pin]
|
|
arrow = Arrow(direction=parse_arrow_str(arrow_str), weight=ArrowWeight.SINGLE)
|
|
|
|
self.mates.append(MatePin(from_pin_obj, to_pin_obj, arrow))
|
|
self.connectors[from_name].activate_pin(
|
|
from_pin, Side.RIGHT, is_connection=False
|
|
)
|
|
self.connectors[to_name].activate_pin(to_pin, Side.LEFT, is_connection=False)
|
|
|
|
def add_mate_component(self, from_name, to_name, arrow_str) -> None:
|
|
arrow = Arrow(direction=parse_arrow_str(arrow_str), weight=ArrowWeight.SINGLE)
|
|
self.mates.append(MateComponent(from_name, to_name, arrow))
|
|
|
|
def populate_bom(self):
|
|
# helper lists
|
|
all_toplevel_items = (
|
|
list(self.connectors.values())
|
|
+ list(self.cables.values())
|
|
+ self.additional_bom_items
|
|
)
|
|
all_subitems = [
|
|
subitem
|
|
for item in all_toplevel_items
|
|
for subitem in item.additional_components
|
|
]
|
|
all_bom_relevant_items = (
|
|
list(self.connectors.values())
|
|
+ [cable for cable in self.cables.values() if cable.category != "bundle"]
|
|
+ [
|
|
wire
|
|
for cable in self.cables.values()
|
|
if cable.category == "bundle"
|
|
for wire in cable.wire_objects.values()
|
|
]
|
|
+ all_subitems
|
|
)
|
|
|
|
def add_to_bom(entry):
|
|
if isinstance(entry, list):
|
|
for e in entry:
|
|
add_to_bom(e)
|
|
return
|
|
|
|
if hash(entry) in self.bom:
|
|
self.bom[hash(entry)] += entry
|
|
else:
|
|
self.bom[hash(entry)] = entry
|
|
|
|
try:
|
|
self.bom[hash(entry)]
|
|
except KeyError:
|
|
raise RuntimeError(
|
|
f"BomEntry's hash is not persitent: h1:{hash(entry)} h2:{hash(entry)}\n\tentry: {entry}\n\titem:{item}"
|
|
)
|
|
|
|
# add items to BOM
|
|
for item in all_toplevel_items:
|
|
if item.ignore_in_bom:
|
|
continue
|
|
add_to_bom(item.bom_entry)
|
|
|
|
# sort BOM by category first, then alphabetically by description within category
|
|
self.bom = dict(
|
|
sorted(
|
|
self.bom.items(),
|
|
key=lambda x: (
|
|
x[1].category,
|
|
x[1].description,
|
|
), # x[0] = key, x[1] = value
|
|
)
|
|
)
|
|
|
|
next_id = len(self.shared_bom) + 1
|
|
# TODO: for each harness, track a (harness_name, qty) pair
|
|
def get_per_harness(v):
|
|
d = {
|
|
"qty": v["qty"],
|
|
}
|
|
return (self.name, d)
|
|
|
|
for key, values in self.bom.items():
|
|
if key in self.shared_bom:
|
|
self.shared_bom[key]["qty"] += values["qty"]
|
|
values["id"] = self.shared_bom[key]["id"]
|
|
else:
|
|
self.shared_bom[key] = values
|
|
self.shared_bom[key]["id"] = next_id
|
|
values["id"] = next_id
|
|
next_id += 1
|
|
|
|
k, v = get_per_harness(values)
|
|
self.shared_bom[key].per_harness[k] = v
|
|
|
|
# print(f'bom length: {len(self.bom)}, shared_bom length: {len(self.shared_bom)}') # for debugging
|
|
|
|
# set BOM IDs within components (for BOM bubbles)
|
|
for item in all_bom_relevant_items:
|
|
if item.ignore_in_bom:
|
|
continue
|
|
if hash(item) not in self.bom:
|
|
continue
|
|
item.id = self.bom[hash(item)].id
|
|
|
|
self.bom = dict(
|
|
sorted(
|
|
self.bom.items(),
|
|
key=lambda x: (x[1].id,),
|
|
)
|
|
)
|
|
# from wireviz.wv_bom import print_bom_table ; print_bom_table(self.bom) # for debugging
|
|
|
|
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:
|
|
# check from and to connectors
|
|
for (name, pin) in zip([from_name, to_name], [from_pin, to_pin]):
|
|
if name is not None and name in self.connectors:
|
|
connector = self.connectors[name]
|
|
# check if provided name is ambiguous
|
|
if pin in connector.pins and pin in connector.pinlabels:
|
|
if connector.pins.index(pin) != connector.pinlabels.index(pin):
|
|
raise Exception(
|
|
f"{name}:{pin} is defined both in pinlabels and pins, "
|
|
"for different pins."
|
|
)
|
|
# TODO: Maybe issue a warning if present in both lists
|
|
# but referencing the same pin?
|
|
if pin in connector.pinlabels:
|
|
if connector.pinlabels.count(pin) > 1:
|
|
raise Exception(f"{name}:{pin} is defined more than once.")
|
|
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.")
|
|
|
|
# check via cable
|
|
if via_name in self.cables:
|
|
cable = self.cables[via_name]
|
|
# check if provided name is ambiguous
|
|
if via_wire in cable.colors and via_wire in cable.wirelabels:
|
|
if cable.colors.index(via_wire) != cable.wirelabels.index(via_wire):
|
|
raise Exception(
|
|
f"{via_name}:{via_wire} is defined both in colors and wirelabels, "
|
|
"for different wires."
|
|
)
|
|
# TODO: Maybe issue a warning if present in both lists
|
|
# but referencing the same wire?
|
|
if via_wire in cable.colors:
|
|
if cable.colors.count(via_wire) > 1:
|
|
raise Exception(
|
|
f"{via_name}:{via_wire} is used for more than one wire."
|
|
)
|
|
# list index starts at 0, wire IDs start at 1
|
|
via_wire = cable.colors.index(via_wire) + 1
|
|
elif via_wire in cable.wirelabels:
|
|
if cable.wirelabels.count(via_wire) > 1:
|
|
raise Exception(
|
|
f"{via_name}:{via_wire} is used for more than one wire."
|
|
)
|
|
via_wire = (
|
|
cable.wirelabels.index(via_wire) + 1
|
|
) # list index starts at 0, wire IDs start at 1
|
|
|
|
# perform the actual connection
|
|
if from_name is not None:
|
|
from_con = self.connectors[from_name]
|
|
from_pin_obj = from_con.pin_objects[from_pin]
|
|
else:
|
|
from_pin_obj = None
|
|
if to_name is not None:
|
|
to_con = self.connectors[to_name]
|
|
to_pin_obj = to_con.pin_objects[to_pin]
|
|
else:
|
|
to_pin_obj = None
|
|
|
|
self.cables[via_name]._connect(from_pin_obj, via_wire, to_pin_obj)
|
|
if from_name in self.connectors:
|
|
self.connectors[from_name].activate_pin(from_pin, Side.RIGHT)
|
|
if to_name in self.connectors:
|
|
self.connectors[to_name].activate_pin(to_pin, Side.LEFT)
|
|
|
|
def create_graph(self) -> Graph:
|
|
dot = Graph()
|
|
set_dot_basics(dot, self.options)
|
|
|
|
for connector in self.connectors.values():
|
|
# generate connector node
|
|
gv_html = gv_node_component(connector)
|
|
gv_html.update_attribs(
|
|
bgcolor=calculate_node_bgcolor(connector, self.options)
|
|
)
|
|
dot.node(
|
|
connector.designator,
|
|
label=f"<\n{gv_html}\n>",
|
|
shape="box",
|
|
style="filled",
|
|
)
|
|
# generate edges for connector loops
|
|
if len(connector.loops) > 0:
|
|
dot.attr("edge", color="#000000")
|
|
loops = gv_connector_loops(connector)
|
|
for head, tail in loops:
|
|
dot.edge(head, tail)
|
|
|
|
# determine if there are double- or triple-colored wires in the harness;
|
|
# if so, pad single-color wires to make all wires of equal thickness
|
|
wire_is_multicolor = [
|
|
len(wire.color) > 1
|
|
for cable in self.cables.values()
|
|
for wire in cable.wire_objects.values()
|
|
]
|
|
if any(wire_is_multicolor):
|
|
wireviz.wv_colors.padding_amount = 3
|
|
else:
|
|
wireviz.wv_colors.padding_amount = 1
|
|
|
|
for cable in self.cables.values():
|
|
# generate cable node
|
|
# TODO: PN info for bundles (per wire)
|
|
gv_html = gv_node_component(cable)
|
|
gv_html.update_attribs(bgcolor=calculate_node_bgcolor(cable, self.options))
|
|
style = "filled,dashed" if cable.category == "bundle" else "filled"
|
|
dot.node(
|
|
cable.designator,
|
|
label=f"<\n{gv_html}\n>",
|
|
shape="box",
|
|
style=style,
|
|
)
|
|
|
|
# generate wire edges between component nodes and cable nodes
|
|
for connection in cable._connections:
|
|
color, l1, l2, r1, r2 = gv_edge_wire(self, cable, connection)
|
|
dot.attr("edge", color=color)
|
|
if not (l1, l2) == (None, None):
|
|
dot.edge(l1, l2)
|
|
if not (r1, r2) == (None, None):
|
|
dot.edge(r1, r2)
|
|
|
|
for mate in self.mates:
|
|
color, dir, code_from, code_to = gv_edge_mate(mate)
|
|
|
|
dot.attr("edge", color=color, style="dashed", dir=dir)
|
|
dot.edge(code_from, code_to)
|
|
|
|
return dot
|
|
|
|
# cache for the GraphViz Graph object
|
|
# do not access directly, use self.graph instead
|
|
_graph = None
|
|
|
|
@property
|
|
def graph(self):
|
|
if not self._graph: # no cached graph exists, generate one
|
|
self._graph = self.create_graph()
|
|
return self._graph # return cached graph
|
|
|
|
@property
|
|
def png(self):
|
|
from io import BytesIO
|
|
|
|
graph = self.graph
|
|
data = BytesIO()
|
|
data.write(graph.pipe(format="png"))
|
|
data.seek(0)
|
|
return data.read()
|
|
|
|
@property
|
|
def svg(self):
|
|
graph = self.graph
|
|
return embed_svg_images(graph.pipe(format="svg").decode("utf-8"), Path.cwd())
|
|
|
|
def output(
|
|
self,
|
|
filename: (str, Path),
|
|
view: bool = False,
|
|
cleanup: bool = True,
|
|
fmt: tuple = ("html", "png", "svg", "tsv"),
|
|
) -> None:
|
|
# graphical output
|
|
graph = self.graph
|
|
|
|
rendered = set()
|
|
for f in fmt:
|
|
if f in ("png", "svg", "html"):
|
|
if f == "html": # if HTML format is specified,
|
|
f = "svg" # generate SVG for embedding into HTML
|
|
# SVG file will be renamed/deleted later
|
|
if f in rendered:
|
|
continue
|
|
graph.format = f
|
|
graph.render(filename=filename, view=view, cleanup=cleanup)
|
|
rendered.add(f)
|
|
# embed images into SVG output
|
|
if "svg" in fmt or "html" in fmt:
|
|
embed_svg_images_file(filename.with_suffix(".svg"))
|
|
# GraphViz output
|
|
if "gv" in fmt:
|
|
graph.save(filename=filename.with_suffix(".gv"))
|
|
# BOM output
|
|
bomlist = bom_list(self.bom)
|
|
if "tsv" in fmt:
|
|
bom_tsv = bom2tsv(bomlist)
|
|
filename.with_suffix(".tsv").open("w").write(bom_tsv)
|
|
if "csv" in fmt:
|
|
# TODO: implement CSV output (preferrably using CSV library)
|
|
print("CSV output is not yet supported")
|
|
# HTML output
|
|
if "html" in fmt:
|
|
generate_html_output(filename, bomlist, self.metadata, self.options)
|
|
# PDF output
|
|
if "pdf" in fmt:
|
|
# TODO: implement PDF output
|
|
print("PDF output is not yet supported")
|
|
# delete SVG if not needed
|
|
if "html" in fmt and not "svg" in fmt:
|
|
# SVG file was just needed to generate HTML
|
|
filename.with_suffix(".svg").unlink()
|