wv_harness: refactor bom management

src/wireviz/wv_harness.py

@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
 
-from collections import defaultdict
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Dict, List
@@ -8,12 +7,11 @@ from typing import Dict, List
 from graphviz import Graph
 
 import wireviz.wv_colors
-from wireviz.wv_bom import BomCategory, bom_list
+from wireviz.wv_bom import bom_list
 from wireviz.wv_dataclasses import (
-    AUTOGENERATED_PREFIX,
-    AdditionalComponent,
     Arrow,
     ArrowWeight,
+    BomCategory,
     Cable,
     Component,
     Connector,
@@ -22,7 +20,6 @@ from wireviz.wv_dataclasses import (
     Metadata,
     Options,
     Side,
-    TopLevelGraphicalComponent,
     Tweak,
 )
 from wireviz.wv_graphviz import (
@@ -44,16 +41,22 @@ class Harness:
     metadata: Metadata
     options: Options
     tweak: Tweak
-    additional_bom_items: List[AdditionalComponent] = field(default_factory=list)
+    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 = defaultdict(dict)
+        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
@@ -63,7 +66,7 @@ class Harness:
         self.cables[designator] = cbl
 
     def add_additional_bom_item(self, item: dict) -> None:
-        new_item = AdditionalComponent(**item)
+        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:
@@ -107,124 +110,79 @@ class Harness:
             + 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:
-            self._add_to_internal_bom(item)  # nested subitems are also handled
+            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].category,
-                    x[0].description,
+                    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"]
-                continue
+            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 not item.bom_hash in self.bom:
+            if hash(item) not in self.bom:
-                print(f"{item}'s hash' not found in BOM dict.")
                 continue
-            item.bom_id = self.bom[item.bom_hash]["id"]
+            item.id = self.bom[hash(item)].id
 
         self.bom = dict(
             sorted(
                 self.bom.items(),
-                key=lambda x: (x[1]["id"],),
+                key=lambda x: (x[1].id,),
             )
         )
-        # print_bom_table(self.bom)  # for debugging
+        # from wireviz.wv_bom import print_bom_table ; print_bom_table(self.bom)  # for debugging
-
-    def _add_to_internal_bom(self, item: Component):
-        if item.ignore_in_bom:
-            return
-
-        def _add(hash, qty, designator=None, category=None):
-            bom_entry = self.bom[hash]
-            # initialize missing fields
-            if not "qty" in bom_entry:
-                bom_entry["qty"] = 0
-            if not "designators" in bom_entry:
-                bom_entry["designators"] = set()
-            # update fields
-            bom_entry["qty"] += qty
-            if designator is None:
-                designator_list = []
-            elif isinstance(designator, list):
-                designator_list = designator
-            else:
-                designator_list = [designator]
-            for des in designator_list:
-                if des and not des.startswith(AUTOGENERATED_PREFIX):
-                    bom_entry["designators"].add(des)
-            bom_entry["category"] = category
-
-        if isinstance(item, TopLevelGraphicalComponent):
-            if isinstance(item, Connector):
-                cat = BomCategory.CONNECTOR
-            elif isinstance(item, Cable):
-                if item.category == "bundle":
-                    cat = BomCategory.WIRE
-                else:
-                    cat = BomCategory.CABLE
-            else:
-                cat = ""
-
-            if item.category == "bundle":
-                for subitem in item.wire_objects.values():
-                    _add(
-                        hash=subitem.bom_hash,
-                        qty=item.bom_qty,  # should be 1
-                        designator=item.designator,  # inherit from parent item
-                        category=cat,
-                    )
-            else:
-                _add(
-                    hash=item.bom_hash,
-                    qty=item.bom_qty,
-                    designator=item.designator,
-                    category=cat,
-                )
-            if item.additional_components:
-                if item.category == "bundle":
-                    pass  # TODO
-                item.compute_qty_multipliers()
-            for comp in item.additional_components:
-                if comp.ignore_in_bom:
-                    continue
-                _add(
-                    hash=comp.bom_hash,
-                    designator=item.designator,
-                    qty=comp.bom_qty,
-                    category=BomCategory.ADDITIONAL_INSIDE,
-                )
-        elif isinstance(item, AdditionalComponent):
-            cat = BomCategory.ADDITIONAL_OUTSIDE
-            _add(
-                hash=item.bom_hash,
-                qty=item.bom_qty,
-                designator=None,
-                category=cat,
-            )
-        else:
-            raise Exception(f"Unknown type of item:\n{item}")
 
     def connect(
         self,