Make connecting things more object-oriented

src/wireviz/DataClasses.py

@@ -2,6 +2,7 @@
 
 from dataclasses import InitVar, dataclass, field
 from enum import Enum, auto
+from itertools import zip_longest
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -146,6 +147,42 @@ class Component:
     pass
 
 
+@dataclass
+class PinClass:
+    index: int
+    id: str
+    label: str
+    color: str
+    parent: str  # designator of parent connector
+
+
+@dataclass
+class WireClass:
+    index: int
+    id: str
+    label: str
+    color: str
+    parent: str  # designator of parent cable/bundle
+    # gauge: Gauge
+    # pn: str
+    # manufacturer: str
+    # mpn: str
+    # supplier: str
+    # spn: str
+
+
+@dataclass
+class ShieldClass(WireClass):
+    pass  # TODO, for wires with multiple shields more shield details, ...
+
+
+@dataclass
+class Connection:
+    from_: PinClass = None
+    via: Union[WireClass, ShieldClass] = None
+    to: PinClass = None
+
+
 @dataclass
 class Connector(Component):
     name: Designator
@@ -173,6 +210,7 @@ class Connector(Component):
     loops: List[List[Pin]] = field(default_factory=list)
     ignore_in_bom: bool = False
     additional_components: List[AdditionalComponent] = field(default_factory=list)
+    pin_objects: List[PinClass] = field(default_factory=list)
 
     @property
     def is_autogenerated(self):
@@ -213,6 +251,23 @@ class Connector(Component):
         if len(self.pins) != len(set(self.pins)):
             raise Exception("Pins are not unique")
 
+        # all checks have passed
+        pin_tuples = zip_longest(
+            self.pins,
+            self.pinlabels,
+            self.pincolors,
+        )
+        for pin_index, (pin_id, pin_label, pin_color) in enumerate(pin_tuples):
+            self.pin_objects.append(
+                PinClass(
+                    index=pin_index,
+                    id=pin_id,
+                    label=pin_label,
+                    color=pin_color,
+                    parent=self.name,
+                )
+            )
+
         if self.show_name is None:
             self.show_name = self.style != "simple" and not self.is_autogenerated
 
@@ -235,6 +290,19 @@ class Connector(Component):
             if isinstance(item, dict):
                 self.additional_components[i] = AdditionalComponent(**item)
 
+    def _check_if_unique_id(self, id):
+        results = [pin for pin in self.pin_objects if pin.id == id]
+        if len(results) == 0:
+            raise Exception(f"Pin ID {id} not found in {self.name}")
+        if len(results) > 1:
+            raise Exception(f"Pin ID {id} found more than once in {self.name}")
+        return True
+
+    def get_pin_by_id(self, id):
+        if self._check_if_unique_id(id):
+            pin = [pin for pin in self.pin_objects if pin.id == id]
+            return pin[0]
+
     def activate_pin(self, pin: Pin, side: Side) -> None:
         self.visible_pins[pin] = True
         if side == Side.LEFT:
@@ -285,6 +353,8 @@ class Cable(Component):
     show_wirenumbers: Optional[bool] = None
     ignore_in_bom: bool = False
     additional_components: List[AdditionalComponent] = field(default_factory=list)
+    connections: List[Connection] = field(default_factory=list)
+    wire_objects: List[WireClass] = field(default_factory=list)
 
     @property
     def is_autogenerated(self):
@@ -354,8 +424,6 @@ class Cable(Component):
         elif self.length_unit is None:
             self.length_unit = "m"
 
-        self.connections = []
-
         if self.wirecount:  # number of wires explicitly defined
             if self.colors:  # use custom color palette (partly or looped if needed)
                 pass
@@ -396,6 +464,36 @@ class Cable(Component):
                 else:
                     raise Exception("lists of part data are only supported for bundles")
 
+        # all checks have passed
+        wire_tuples = zip_longest(
+            # TODO: self.wire_ids
+            self.colors,
+            self.wirelabels,
+        )
+        for wire_index, (wire_color, wire_label) in enumerate(wire_tuples):
+            self.wire_objects.append(
+                WireClass(
+                    index=wire_index,  # TODO: wire_id
+                    id=wire_index + 1,  # TODO: wire_id
+                    label=wire_label,
+                    color=wire_color,
+                    parent=self.name,
+                )
+            )
+
+        if self.shield:
+            index_offset = len(self.wire_objects)
+            # TODO: add support for multiple shields
+            self.wire_objects.append(
+                ShieldClass(
+                    index=index_offset,
+                    id="s",
+                    label="Shield",
+                    color=self.shield if isinstance(self.shield, str) else None,
+                    parent=self.name,
+                )
+            )
+
         if self.show_name is None:
             self.show_name = not self.is_autogenerated
 
@@ -407,25 +505,19 @@ class Cable(Component):
             if isinstance(item, dict):
                 self.additional_components[i] = AdditionalComponent(**item)
 
-    # The *_pin arguments accept a tuple, but it seems not in use with the current code.
+    def get_wire_by_id(self, id):
-    def connect(
+        wire = [wire for wire in self.wire_objects if wire.id == id]
-        self,
+        if len(wire) == 0:
-        from_name: Optional[Designator],
+            raise Exception(f"Wire ID {id} not found in {self.name}")
-        from_pin: NoneOrMorePinIndices,
+        if len(wire) > 1:
-        via_wire: OneOrMoreWires,
+            raise Exception(f"Wire ID {id} found more than once in {self.name}")
-        to_name: Optional[Designator],
+        return wire[0]
-        to_pin: NoneOrMorePinIndices,
-    ) -> None:
 
-        from_pin = int2tuple(from_pin)
+    def connect(
-        via_wire = int2tuple(via_wire)
+        self, from_pin_obj: [PinClass], via_wire_id: str, to_pin_obj: [PinClass]
-        to_pin = int2tuple(to_pin)
+    ) -> None:
-        if len(from_pin) != len(to_pin):
+        via_wire_obj = self.get_wire_by_id(via_wire_id)
-            raise Exception("from_pin must have the same number of elements as to_pin")
+        self.connections.append(Connection(from_pin_obj, via_wire_obj, to_pin_obj))
-        for i, _ in enumerate(from_pin):
-            self.connections.append(
-                Connection(from_name, from_pin[i], via_wire[i], to_name, to_pin[i])
-            )
 
     def get_qty_multiplier(self, qty_multiplier: Optional[CableMultiplier]) -> float:
         if not qty_multiplier:
@@ -444,15 +536,6 @@ class Cable(Component):
             )
 
 
-@dataclass
-class Connection:
-    from_name: Optional[Designator]
-    from_pin: Optional[Pin]
-    via_port: Wire
-    to_name: Optional[Designator]
-    to_pin: Optional[Pin]
-
-
 @dataclass
 class MatePin:
     from_name: Designator

src/wireviz/Harness.py

@@ -140,7 +140,18 @@ class Harness:
                 )  # list index starts at 0, wire IDs start at 1
 
         # perform the actual connection
-        self.cables[via_name].connect(from_name, from_pin, via_wire, to_name, to_pin)
+        if from_name is not None:
+            from_con = self.connectors[from_name]
+            from_pin_obj = from_con.get_pin_by_id(from_pin)
+        else:
+            from_pin_obj = None
+        if to_name is not None:
+            to_con = self.connectors[to_name]
+            to_pin_obj = to_con.get_pin_by_id(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:

src/wireviz/wv_gv_html.py

@@ -5,7 +5,16 @@ from itertools import zip_longest
 from typing import Any, List, Optional, Union
 
 from wireviz import APP_NAME, APP_URL, __version__
-from wireviz.DataClasses import Cable, Color, Component, Connector, Options
+from wireviz.DataClasses import (
+    Cable,
+    Color,
+    Component,
+    Connection,
+    Connector,
+    Options,
+    ShieldClass,
+    WireClass,
+)
 from wireviz.wv_colors import get_color_hex, translate_color
 from wireviz.wv_helper import pn_info_string, remove_links
 from wireviz.wv_table_util import *  # TODO: explicitly import each needed tag later
@@ -154,6 +163,7 @@ def gv_pin_table(component) -> Table:
 
 
 def gv_pin_row(pin_index, pin_name, pin_label, pin_color, connector) -> Tr:
+    # ports in GraphViz are 1-indexed for more natural maping to pin/wire numbers
     cell_pin_left = Td(pin_name, port=f"p{pin_index+1}l")
     cell_pin_label = Td(pin_label, delete_if_empty=True)
     cell_pin_right = Td(pin_name, port=f"p{pin_index+1}r")
@@ -187,61 +197,56 @@ def gv_conductor_table(cable, harness_options) -> Table:
     rows = []
     rows.append(Tr(Td(" ")))  # spacer row on top
 
-    for i, (connection_color, wirelabel) in enumerate(
+    inserted_break_inbetween = False
-        zip_longest(cable.colors, cable.wirelabels), 1
+    for wire in cable.wire_objects:
-    ):
+
+        # insert blank space between wires and shields
+        if isinstance(wire, ShieldClass) and not inserted_break_inbetween:
+            rows.append(Tr(Td(" ")))  # spacer row between wires and shields
+            inserted_break_inbetween = True
 
         # row above the wire
         wireinfo = []
-        if cable.show_wirenumbers:
+        if cable.show_wirenumbers and not isinstance(wire, ShieldClass):
-            wireinfo.append(str(i))
+            wireinfo.append(str(wire.id))
-        colorstr = translate_color(connection_color, harness_options.color_mode)
+        wireinfo.append(translate_color(wire.color, harness_options.color_mode))
-        if colorstr:
+        wireinfo.append(wire.label)
-            wireinfo.append(colorstr)
+
-        if cable.wirelabels:
+        ins, outs = [], []
-            wireinfo.append(wirelabel if wirelabel is not None else "")
+        for conn in cable.connections:
+            if conn.via.id == wire.id:
+                if conn.from_ is not None:
+                    from_label = f":{conn.from_.label}" if conn.from_.label else ""
+                    ins.append(f"{conn.from_.parent}:{conn.from_.id}{from_label}")
+                if conn.to is not None:
+                    to_label = f":{conn.to.label}" if conn.to.label else ""
+                    outs.append(f"{conn.to.parent}:{conn.to.id}{to_label}")
 
         cells_above = [
-            Td(f"<!-- {i}_in -->"),
+            Td(", ".join(ins)),
-            Td(":".join(wireinfo)),
+            Td(":".join([wi for wi in wireinfo if wi is not None])),
-            Td(f"<!-- {i}_out -->"),
+            Td(", ".join(outs)),
         ]
         rows.append(Tr(cells_above))
 
         # the wire itself
-        color_list = (
+        rows.append(Tr(gv_wire_cell(wire, padding=harness_options._pad)))
-            ["#000000"]
-            + get_color_hex(connection_color, pad=harness_options._pad)
-            + ["#000000"]
-        )
-        rows.append(Tr(gv_wire_cell(i, color_list)))
 
         # row below the wire
         # TODO: PN stuff for bundles
         # wire_pn_stuff() see below
 
-    if cable.shield:
-        rows.append(Tr(Td("&nbsp;")))  # spacer between wires and shield
-        # row above the shield
-        cells_above = [
-            Td("<!-- s_in -->"),
-            Td("Shield"),
-            Td("<!-- s_out -->"),
-        ]
-        rows.append(Tr(cells_above))
-        # thw shield itself
-        if isinstance(cable.shield, str):
-            color_list = ["#000000"] + get_color_hex(cable.shield) + ["#000000"]
-        else:
-            color_list = ["#000000"]
-        rows.append(Tr(gv_wire_cell("s", color_list)))
-
     rows.append(Tr(Td("&nbsp;")))  # spacer row on bottom
     tbl = Table(rows, border=0, cellspacing=0, cellborder=0)
     return tbl
 
 
-def gv_wire_cell(index, color_list) -> Td:
+def gv_wire_cell(wire: Union[WireClass, ShieldClass], padding) -> Td:
+    if wire.color:
+        color_list = ["#000000"] + get_color_hex(wire.color, pad=padding) + ["#000000"]
+    else:
+        color_list = ["#000000"]
+
     wire_inner_rows = []
     for j, bgcolor in enumerate(color_list[::-1]):
         wire_inner_cell_attribs = {
@@ -257,9 +262,10 @@ def gv_wire_cell(index, color_list) -> Td:
         "colspan": 3,
         "border": 0,
         "cellspacing": 0,
-        "port": f"w{index}",
+        "port": f"w{wire.index+1}",
         "height": 2 * len(color_list),
     }
+    # ports in GraphViz are 1-indexed for more natural maping to pin/wire numbers
     wire_outer_cell = Td(wire_inner_table, **wire_outer_cell_attribs)
 
     return wire_outer_cell
@@ -308,69 +314,39 @@ def wire_pn_stuff():
 
 
 def gv_edge_wire(harness, cable, connection) -> (str, str, str):
-    if isinstance(connection.via_port, int):
+    if connection.via.color:
         # check if it's an actual wire and not a shield
-        wire_color = get_color_hex(
+        wire_color = get_color_hex(connection.via.color, pad=harness.options._pad)
-            cable.colors[connection.via_port - 1], pad=harness.options._pad
-        )
         color = ":".join(["#000000"] + wire_color + ["#000000"])
     else:  # it's a shield connection
         # shield is shown with specified color and black borders, or as a thin black wire otherwise
-        if isinstance(cable.shield, str):
+        if connection.via.color:
-            shield_color_hex = get_color_hex(cable.shield)[0]
+            shield_color_hex = get_color_hex(connection.via.color)[0]
             shield_color_str = ":".join(["#000000", shield_color_hex, "#000000"])
         else:
             shield_color_str = "#000000"
         color = shield_color_str
-    if connection.from_pin is not None:  # connect to left
-        from_connector = harness.connectors[connection.from_name]
-        from_pin_index = from_connector.pins.index(connection.from_pin)
-        from_port_str = (
-            f":p{from_pin_index+1}r" if from_connector.style != "simple" else ""
-        )
-        code_left_1 = f"{connection.from_name}{from_port_str}:e"
-        code_left_2 = f"{cable.name}:w{connection.via_port}:w"
-        # dot.edge(code_left_1, code_left_2)
-        if from_connector.show_name:
-            from_info = [
-                str(connection.from_name),
-                str(connection.from_pin),
-            ]
-            if from_connector.pinlabels:
-                pinlabel = from_connector.pinlabels[from_pin_index]
-                if pinlabel != "":
-                    from_info.append(pinlabel)
-            from_string = ":".join(from_info)
 
-        else:
+    if connection.from_ is not None:  # connect to left
-            from_string = ""
+        from_port_str = (
-        # html = [
+            f":p{connection.from_.index+1}r"
-        #     row.replace(f"<!-- {connection.via_port}_in -->", from_string)
+            if harness.connectors[connection.from_.parent].style != "simple"
-        #     for row in html
+            else ""
-        # ]
+        )
+        code_left_1 = f"{connection.from_.parent}{from_port_str}:e"
+        code_left_2 = f"{connection.via.parent}:w{connection.via.index+1}:w"
+        # ports in GraphViz are 1-indexed for more natural maping to pin/wire numbers
     else:
         code_left_1, code_left_2 = None, None
 
-    if connection.to_pin is not None:  # connect to right
+    if connection.to is not None:  # connect to right
-        to_connector = harness.connectors[connection.to_name]
+        to_port_str = (
-        to_pin_index = to_connector.pins.index(connection.to_pin)
+            f":p{connection.to.index+1}l"
-        to_port_str = f":p{to_pin_index+1}l" if to_connector.style != "simple" else ""
+            if harness.connectors[connection.from_.parent].style != "simple"
-        code_right_1 = f"{cable.name}:w{connection.via_port}:e"
+            else ""
-        code_right_2 = f"{connection.to_name}{to_port_str}:w"
+        )
-        # dot.edge(code_right_1, code_right_2)
+        code_right_1 = f"{connection.via.parent}:w{connection.via.index+1}:e"
-        if to_connector.show_name:
+        code_right_2 = f"{connection.to.parent}{to_port_str}:w"
-            to_info = [str(connection.to_name), str(connection.to_pin)]
-            if to_connector.pinlabels:
-                pinlabel = to_connector.pinlabels[to_pin_index]
-                if pinlabel != "":
-                    to_info.append(pinlabel)
-            to_string = ":".join(to_info)
-        else:
-            to_string = ""
-        # html = [
-        #     row.replace(f"<!-- {connection.via_port}_out -->", to_string)
-        #     for row in html
-        # ]
     else:
         code_right_1, code_right_2 = None, None