From e8bfbd352b350adf7f0040716916a4ea549eba62 Mon Sep 17 00:00:00 2001
From: Daniel Rojas <github@danielrojas.net>
Date: Fri, 23 Oct 2020 18:39:41 +0200
Subject: [PATCH] Add option to define connector pin colors (#141)

---
 src/wireviz/DataClasses.py |  4 ++++
 src/wireviz/Harness.py     | 12 +++++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/wireviz/DataClasses.py b/src/wireviz/DataClasses.py
index e813fbe3..15621287 100644
--- a/src/wireviz/DataClasses.py
+++ b/src/wireviz/DataClasses.py
@@ -78,6 +78,7 @@ class Connector:
     notes: Optional[str] = None
     pinlabels: List[Any] = field(default_factory=list)
     pins: List[Any] = field(default_factory=list)
+    pincolors: List[Any] = field(default_factory=list)
     color: Optional[str] = None
     show_name: bool = None
     show_pincount: bool = None
@@ -122,6 +123,9 @@ def __post_init__(self):
         if len(self.pins) != len(set(self.pins)):
             raise Exception('Pins are not unique')
 
+        if self.pincolors:
+            self.pincolors.extend(["--"] * (len(self.pins) - len(self.pincolors)))  # autofill missing pincolors as 'no color'
+
         if self.show_name is None:
             self.show_name = not self.autogenerate # hide auto-generated designators by default
 
diff --git a/src/wireviz/Harness.py b/src/wireviz/Harness.py
index 835ed651..b167a159 100644
--- a/src/wireviz/Harness.py
+++ b/src/wireviz/Harness.py
@@ -110,7 +110,9 @@ def create_graph(self) -> Graph:
                 pinhtml = []
                 pinhtml.append('<table border="0" cellspacing="0" cellpadding="3" cellborder="1">')
 
-                for pin, pinlabel in zip(connector.pins, connector.pinlabels):
+                for pin, pinlabel, pincolor in zip(connector.pins,
+                                                   connector.pinlabels,
+                                                   connector.pincolors if connector.pincolors else [None] * len(connector.pins)):
                     if connector.hide_disconnected_pins and not connector.visible_pins.get(pin, False):
                         continue
                     pinhtml.append('   <tr>')
@@ -118,6 +120,14 @@ def create_graph(self) -> Graph:
                         pinhtml.append(f'    <td port="p{pin}l">{pin}</td>')
                     if pinlabel:
                         pinhtml.append(f'    <td>{pinlabel}</td>')
+                    if connector.pincolors:
+                        if pincolor in wv_colors._color_hex.keys():
+                            pinhtml.append(f'<td sides="tbl">{pincolor}</td>')
+                            pinhtml.append(f'<td sides="tbr"><table border="0" cellborder="1"><tr><td bgcolor="{wv_colors.translate_color(pincolor, "HEX")}" width="8" height="8" fixedsize="true"></td></tr></table></td>')
+                        else:
+                            pinhtml.append(f'<td sides="tbl"></td>')
+                            pinhtml.append(f'<td sides="tbr"></td>')
+
                     if connector.ports_right:
                         pinhtml.append(f'    <td port="p{pin}r">{pin}</td>')
                     pinhtml.append('   </tr>')