ee.bom: Moving it a single file, no need for more.

1 files changed, 124 insertions, 0 deletions

diff --git a/src/ee/bom.py b/src/ee/bom.py
new file mode 100644
index 0000000..f8cca71
--- /dev/null
+++ b/src/ee/bom.py
@@ -0,0 +1,124 @@
+import os.path
+import pydoc
+from pathlib import Path
+from typing import List, MutableMapping, Optional
+
+from ee import EeException
+from ee.db import ObjDb
+from ee.part import PartDb, load_db, save_db, Part
+from ee.project import Project, report, SupplierDescriptor
+from ee.xml import types
+
+__all__ = ["create_bom"]
+
+
+class BomPart(object):
+    def __init__(self, part: types.Part):
+        self.part = Part(part)
+        ref = self.part.get_only_schematic_reference()
+        self.ref = ref.referenceProp if ref else None
+        self.available_from: MutableMapping[str, Part] = {}
+        self.selected_part = None
+
+
+def make_report(out_file, unresolved_parts, bom_parts: ObjDb[BomPart], supplier_parts: ObjDb[Path]):
+    kwargs = {
+        "bom_parts": bom_parts,
+        "supplier_parts": supplier_parts,
+        "unresolved_parts": unresolved_parts,
+    }
+    report.save_report("ee", "bom.rst.j2", out_file, **kwargs)
+
+
+def default_strategy(x):
+    return x
+
+
+def create_bom(project: Project, schematic_path: Path, out_path: Path, part_dbs: List[Path],
+               fail_on_missing_parts: bool, strategy_name: Optional[str]):
+    strategy = default_strategy
+    if strategy_name:
+        strategy = pydoc.locate(strategy_name)
+        if not callable(strategy):
+            raise EeException("Not a callable: {}, is a {}".format(strategy_name, type(strategy)))
+
+    supplier_parts = ObjDb[Part]()
+    supplier_parts.add_unique_index("uri", lambda p: p.uri)
+    supplier_parts.add_index("spn", lambda p: [ref.valueProp for ref in p.get_spns()],
+                             multiple=True)
+    supplier_pn_idx = supplier_parts.add_multi_index("supplier,pn", lambda p: [
+        (p.supplier, ref.valueProp) for ref in p.get_mpns()], multiple=True)
+    supplier_spn_idx = supplier_parts.add_multi_index("supplier,spn", lambda p: [
+        (p.supplier, ref.valueProp) for ref in p.get_spns()], multiple=True)
+
+    suppliers: List[SupplierDescriptor] = [project.get_supplier_by_key(path.parent.name) for path in part_dbs]
+    for path in part_dbs:
+        for xml in load_db(path).iterparts():
+            if not xml.supplierProp:
+                continue
+            supplier_parts.add(Part(xml))
+
+    sch_db = load_db(schematic_path)
+    bom_parts: ObjDb[BomPart] = ObjDb[BomPart]()
+    bom_parts.add_multi_index("supplier,pn", lambda op: [
+        (op.part.supplierProp, ref.valueProp) for ref in
+        op.part.get_mpns()] if op.part.supplier else None, multiple=True)
+
+    for sch_part in sch_db.iterparts():
+        part = BomPart(sch_part)
+        part.part = strategy(part.part)
+        if part.part is None:
+            continue
+
+        bom_parts.add(part)
+
+    for bom_part in bom_parts:
+        sch_part_numbers = [pn.valueProp for pn in bom_part.part.get_mpns()]
+        sch_supplier_part_numbers = [spn.valueProp for spn in bom_part.part.get_spns()]
+
+        for supplier in suppliers:
+            pns = supplier_pn_idx.get(supplier.uri)
+            spns = supplier_spn_idx.get(supplier.uri)
+
+            for sch_pn in sch_part_numbers:
+                for supplier_part in pns.get(sch_pn, []):
+                    bom_part.available_from[supplier_part.uri] = supplier_part
+            for sch_spn in sch_supplier_part_numbers:
+                for supplier_part in spns.get(sch_spn, []):
+                    bom_part.available_from[supplier_part.uri] = supplier_part
+
+    unresolved_parts = []
+    for bom_part in bom_parts:
+        af = bom_part.available_from
+        if len(af) == 0:
+            unresolved_parts.append(bom_part)
+        elif len(af) == 1:
+            bom_part.selected_part = next(iter(af.values()))
+        else:
+            raise EeException("unimplemented: part ({}) available from multiple suppliers: {}".
+                              format(bom_part.ref, ",".join(af.keys())))
+
+    bom_parts.add_index("uri", lambda op: op.selected_part.uri if op.selected_part else None)
+    bom_parts.add_multi_index("supplier,part", lambda op: (
+        op.selected_part.supplier, op.selected_part.uri) if op.selected_part else None)
+
+    if len(unresolved_parts) and fail_on_missing_parts:
+        raise EeException("The bom has parts that can't be found from any supplier")
+
+    out_file = project.report_dir / (os.path.splitext(out_path.name)[0] + ".rst")
+    make_report(out_file, unresolved_parts, bom_parts, supplier_parts)
+
+    out_parts = PartDb()
+    for bom_part in bom_parts:
+        if not bom_part.selected_part:
+            continue
+
+        supplier_part = bom_part.selected_part
+
+        part = Part(types.Part(supplier=supplier_part.supplier))
+        part.add_schematic_reference(bom_part.part.get_exactly_one_schematic_reference().referenceProp)
+        part.add_part_reference(supplier_part.uri)
+
+        out_parts.add_entry(part, True)
+
+    save_db(out_path, out_parts)