Make mates object-oriented

src/wireviz/DataClasses.py

@@ -42,10 +42,18 @@ MetadataKeys = PlainText  # Literal['title', 'description', 'notes', ...]
 
 
 Side = Enum("Side", "LEFT RIGHT")
+ArrowDirection = Enum("ArrowDirection", "NONE BACK FORWARD BOTH")
+ArrowWeight = Enum("ArrowWeight", "SINGLE DOUBLE")
 
 AUTOGENERATED_PREFIX = "AUTOGENERATED_"
 
 
+@dataclass
+class Arrow:
+    direction: ArrowDirection
+    weight: ArrowWeight
+
+
 class Metadata(dict):
     pass
 
@@ -154,6 +162,16 @@ class PinClass:
     label: str
     color: str
     parent: str  # designator of parent connector
+    _anonymous: bool = False  # true for pins on autogenerated connectors
+    _simple: bool = False  # true for simple connector
+
+    def __str__(self):
+        snippets = [  # use str() for each in case they are int or other non-str
+            str(self.parent) if not self._anonymous else "",
+            str(self.id) if not self._anonymous and not self._simple else "",
+            str(self.label) if self.label else "",
+        ]
+        return ":".join([snip for snip in snippets if snip != ""])
 
 
 @dataclass
@@ -265,6 +283,8 @@ class Connector(Component):
                     label=pin_label,
                     color=pin_color,
                     parent=self.name,
+                    _anonymous=self.is_autogenerated,
+                    _simple=self.style == "simple",
                 )
             )
 
@@ -538,15 +558,13 @@ class Cable(Component):
 
 @dataclass
 class MatePin:
-    from_name: Designator
+    from_: PinClass
-    from_pin: Pin
+    to: PinClass
-    to_name: Designator
+    arrow: Arrow
-    to_pin: Pin
-    shape: str
 
 
 @dataclass
 class MateComponent:
-    from_name: Designator
+    from_: str  # Designator
-    to_name: Designator
+    to: str  # Designator
-    shape: str
+    arrow: Arrow

src/wireviz/Harness.py

@@ -11,6 +11,9 @@ from graphviz import Graph
 
 from wireviz import APP_NAME, APP_URL, __version__, wv_colors
 from wireviz.DataClasses import (
+    Arrow,
+    ArrowDirection,
+    ArrowWeight,
     Cable,
     Connector,
     MateComponent,
@@ -34,12 +37,14 @@ from wireviz.wv_bom import (
 from wireviz.wv_colors import get_color_hex, translate_color
 from wireviz.wv_gv_html import (
     apply_dot_tweaks,
+    calculate_node_bgcolor,
     gv_connector_loops,
+    gv_edge_mate,
     gv_edge_wire,
     gv_node_component,
     html_line_breaks,
+    parse_arrow_str,
     remove_links,
-    calculate_node_bgcolor,
     set_dot_basics,
 )
 from wireviz.wv_helper import (
@@ -71,13 +76,20 @@ class Harness:
     def add_cable(self, name: str, *args, **kwargs) -> None:
         self.cables[name] = Cable(name, *args, **kwargs)
 
-    def add_mate_pin(self, from_name, from_pin, to_name, to_pin, arrow_type) -> None:
+    def add_mate_pin(self, from_name, from_pin, to_name, to_pin, arrow_str) -> None:
-        self.mates.append(MatePin(from_name, from_pin, to_name, to_pin, arrow_type))
+        from_con = self.connectors[from_name]
+        from_pin_obj = from_con.get_pin_by_id(from_pin)
+        to_con = self.connectors[to_name]
+        to_pin_obj = to_con.get_pin_by_id(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)
         self.connectors[to_name].activate_pin(to_pin, Side.LEFT)
 
-    def add_mate_component(self, from_name, to_name, arrow_type) -> None:
+    def add_mate_component(self, from_name, to_name, arrow_str) -> None:
-        self.mates.append(MateComponent(from_name, to_name, arrow_type))
+        arrow = Arrow(direction=parse_arrow_str(arrow_str), weight=ArrowWeight.SINGLE)
+        self.mates.append(MateComponent(from_name, to_name, arrow))
 
     def add_bom_item(self, item: dict) -> None:
         self.additional_bom_items.append(item)
@@ -167,7 +179,11 @@ class Harness:
             gv_html = gv_node_component(connector, self.options)
             bgcolor = calculate_node_bgcolor(connector, self.options)
             dot.node(
-                connector.name, label=f"<\n{gv_html}\n>", bgcolor=bgcolor, shape="box", style="filled"
+                connector.name,
+                label=f"<\n{gv_html}\n>",
+                bgcolor=bgcolor,
+                shape="box",
+                style="filled",
             )
             # generate edges for connector loops
             if len(connector.loops) > 0:
@@ -192,7 +208,15 @@ class Harness:
             gv_html = gv_node_component(cable, self.options)
             bgcolor = calculate_node_bgcolor(cable, self.options)
             style = "filled,dashed" if cable.category == "bundle" else "filled"
-            dot.node(cable.name, label=f"<\n{gv_html}\n>", bgcolor=bgcolor, shape="box", style=style)
+            dot.node(
+                cable.name,
+                label=f"<\n{gv_html}\n>",
+                bgcolor=bgcolor,
+                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)
@@ -201,53 +225,10 @@ class Harness:
                 if not (r1, r2) == (None, None):
                     dot.edge(r1, r2)
 
-
         apply_dot_tweaks(dot, self.tweak)
 
         for mate in self.mates:
-            if mate.shape[0] == "<" and mate.shape[-1] == ">":
+            color, dir, code_from, code_to = gv_edge_mate(mate)
-                dir = "both"
-            elif mate.shape[0] == "<":
-                dir = "back"
-            elif mate.shape[-1] == ">":
-                dir = "forward"
-            else:
-                dir = "none"
-
-            if isinstance(mate, MatePin):
-                color = "#000000"
-            elif isinstance(mate, MateComponent):
-                color = "#000000:#000000"
-            else:
-                raise Exception(f"{mate} is an unknown mate")
-
-            from_connector = self.connectors[mate.from_name]
-            if (
-                isinstance(mate, MatePin)
-                and self.connectors[mate.from_name].style != "simple"
-            ):
-                from_pin_index = from_connector.pins.index(mate.from_pin)
-                from_port_str = f":p{from_pin_index+1}r"
-            else:  # MateComponent or style == 'simple'
-                from_port_str = ""
-
-            to_connector = self.connectors[mate.to_name]
-            if (
-                isinstance(mate, MatePin)
-                and self.connectors[mate.to_name].style != "simple"
-            ):
-                to_pin_index = to_connector.pins.index(mate.to_pin)
-                to_port_str = (
-                    f":p{to_pin_index+1}l"
-                    if isinstance(mate, MatePin)
-                    and self.connectors[mate.to_name].style != "simple"
-                    else ""
-                )
-            else:  # MateComponent or style == 'simple'
-                to_port_str = ""
-            code_from = f"{mate.from_name}{from_port_str}:e"
-            to_connector = self.connectors[mate.to_name]
-            code_to = f"{mate.to_name}{to_port_str}:w"
 
             dot.attr("edge", color=color, style="dashed", dir=dir)
             dot.edge(code_from, code_to)

src/wireviz/wv_gv_html.py

@@ -6,11 +6,16 @@ from typing import Any, List, Optional, Union
 
 from wireviz import APP_NAME, APP_URL, __version__
 from wireviz.DataClasses import (
+    Arrow,
+    ArrowDirection,
+    ArrowWeight,
     Cable,
     Color,
     Component,
     Connection,
     Connector,
+    MateComponent,
+    MatePin,
     Options,
     ShieldClass,
     WireClass,
@@ -41,9 +46,16 @@ def gv_node_component(
 
     line_pn = part_number_str_list(component)
 
+    is_simple_connector = (
+        isinstance(component, Connector) and component.style == "simple"
+    )
+
     if isinstance(component, Connector):
         line_info = [
+            Td(
                 html_line_breaks(component.type),
+                port="p1l" if is_simple_connector else None,
+            ),
             html_line_breaks(component.subtype),
             f"{component.pincount}-pin" if component.show_pincount else None,
             translate_color(component.color, harness_options.color_mode),
@@ -87,6 +99,7 @@ def gv_node_component(
     ]
 
     tbl = nested_table(lines)
+    tbl.update_attribs(port="p1r" if is_simple_connector else None)
 
     return tbl
 
@@ -98,7 +111,11 @@ def calculate_node_bgcolor(component, harness_options):
         return translate_color(component.bgcolor, "HEX")
     elif isinstance(component, Connector) and harness_options.bgcolor_connector:
         return translate_color(harness_options.bgcolor_connector, "HEX")
-    elif isinstance(component, Cable) and component.category == "bundle" and harness_options.bgcolor_bundle:
+    elif (
+        isinstance(component, Cable)
+        and component.category == "bundle"
+        and harness_options.bgcolor_bundle
+    ):
         return translate_color(harness_options.bgcolor_bundle, "HEX")
     elif isinstance(component, Cable) and harness_options.bgcolor_cable:
         return translate_color(harness_options.bgcolor_cable, "HEX")
@@ -222,16 +239,14 @@ def gv_conductor_table(cable, harness_options) -> Table:
         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(str(conn.from_))
-                    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(str(conn.to))
-                    outs.append(f"{conn.to.parent}:{conn.to.id}{to_label}")
 
         cells_above = [
-            Td(", ".join(ins)),
+            Td(", ".join(ins), align="left"),
-            Td(":".join([wi for wi in wireinfo if wi is not None])),
+            Td(":".join([wi for wi in wireinfo if wi is not None and wi != ""])),
-            Td(", ".join(outs)),
+            Td(", ".join(outs), align="right"),
         ]
         rows.append(Tr(cells_above))
 
@@ -349,7 +364,7 @@ def gv_edge_wire(harness, cable, connection) -> (str, str, str):
     if connection.to is not None:  # connect to right
         to_port_str = (
             f":p{connection.to.index+1}l"
-            if harness.connectors[connection.from_.parent].style != "simple"
+            if harness.connectors[connection.to.parent].style != "simple"
             else ""
         )
         code_right_1 = f"{connection.via.parent}:w{connection.via.index+1}:e"
@@ -360,6 +375,46 @@ def gv_edge_wire(harness, cable, connection) -> (str, str, str):
     return color, code_left_1, code_left_2, code_right_1, code_right_2
 
 
+def parse_arrow_str(inp: str) -> ArrowDirection:
+    if inp[0] == "<" and inp[-1] == ">":
+        return ArrowDirection.BOTH
+    elif inp[0] == "<":
+        return ArrowDirection.BACK
+    elif inp[-1] == ">":
+        return ArrowDirection.FORWARD
+    else:
+        return ArrowDirection.NONE
+
+
+def gv_edge_mate(mate) -> (str, str, str, str):
+    if mate.arrow.weight == ArrowWeight.SINGLE:
+        color = "#000000"
+    elif mate.arrow.weight == ArrowWeight.DOUBLE:
+        color = "#000000:#000000"
+
+    dir = mate.arrow.direction.name.lower()
+
+    if isinstance(mate, MatePin):
+        from_pin_index = mate.from_.index
+        from_port_str = f":p{from_pin_index+1}r"
+        from_designator = mate.from_.parent
+        to_pin_index = mate.to.index
+        to_port_str = f":p{to_pin_index+1}l"
+        to_designator = mate.to.parent
+    elif isinstance(mate, MateComponent):
+        from_designator = mate.from_
+        from_port_str = ""
+        to_designator = mate.to
+        to_port_str = ""
+    else:
+        raise Exception(f"Unknown type of mate:\n{mate}")
+
+    code_from = f"{from_designator}{from_port_str}:e"
+    code_to = f"{to_designator}{to_port_str}:w"
+
+    return color, dir, code_from, code_to
+
+
 def colored_cell(contents, bgcolor) -> Td:
     return Td(contents, bgcolor=translate_color(bgcolor, "HEX"))