1 files changed, 151 insertions, 0 deletions
diff --git a/src/ee/kicad/make_bom.py b/src/ee/kicad/make_bom.py
new file mode 100644
index 0000000..3fa102f
--- /dev/null
+++ b/src/ee/kicad/make_bom.py
@@ -0,0 +1,151 @@
+import sys
+from typing import Optional, List, Callable, Mapping
+from xml.dom import minidom
+from xml.etree import ElementTree
+
+from ee import EeException
+from ee.kicad.model import Component
+from ee.kicad.read_schematic import read_schematics
+from ee.kicad.to_bom import to_bom, to_bom_xml
+from ee.tools import mk_parents
+from ee.xml import bomFile
+from ee.bom import part_type_uris
+
+__all__ = [
+    "StrategyCallable",
+    "MakeBomStrategy",
+    "apply_strategies",
+    "mpn_strategy",
+    "part_type_strategy",
+    "dpn_strategy_factory",
+    "make_bom",
+]
+
+StrategyCallable = Callable[[Component, bomFile.Part], Optional[bomFile.Part]]
+
+
+def apply_strategies(c: Component, part: bomFile.Part, strategies: List[StrategyCallable]):
+    for strategy in strategies:
+        part = strategy(c, part)
+
+        if part is None:
+            return
+
+        if not isinstance(part, bomFile.Part):
+            raise EeException("Values returned from strategy must be a bomFile.Part, got {}".format(type(part)))
+
+    return part
+
+
+def part_type_strategy(component: Component, part: bomFile.Part) -> bomFile.Part:
+    fp = component.footprint
+    if fp is None:
+        return part
+
+    lib, part_name = fp.split(":")
+
+    if lib == "Capacitor_SMD":
+        part.part_typeProp = part_type_uris.CAPACITOR
+    elif lib == "Resistor_SMD":
+        part.part_typeProp = part_type_uris.RESISTOR
+    elif lib == "Diode_SMD":
+        part.part_typeProp = part_type_uris.DIODE
+    elif lib == "Inductor_SMD":
+        part.part_typeProp = part_type_uris.INDUCTOR
+    elif lib == "Crystal":
+        part.part_typeProp = part_type_uris.CRYSTAL
+
+    return part
+
+
+def mpn_strategy(component: Component, part: bomFile.Part) -> bomFile.Part:
+    mpn = component.get_field("mpn")
+    if mpn is not None:
+        pn = bomFile.PartNumber(value=mpn.value)
+        part.part_numbersProp.add_part_number(pn)
+
+    return part
+
+
+def dpn_strategy_factory(dpn_mappings: Mapping[str, str]) -> StrategyCallable:
+    def dpn_strategy(component: Component, part: bomFile.Part) -> bomFile.Part:
+        for field_name, distributor in dpn_mappings:
+            s = component.get_field(field_name)
+            if s is not None:
+                pn = bomFile.PartNumber(value=s, distributor=distributor)
+                part.part_numbersProp.add_part_number(pn)
+
+        return part
+
+    return dpn_strategy
+
+
+class MakeBomStrategy():
+    def __init__(self, dpn_mappings: Mapping[str, str] = None):
+        self.dpn_mappings = dpn_mappings or {}
+        self.default_strategies = [
+            mpn_strategy,
+            part_type_strategy,
+            dpn_strategy_factory(self.dpn_mappings),
+        ]
+
+    def process_part(self, component: Component, part: bomFile.Part):
+        return self.default_process_part(component, part)
+
+    def default_process_part(self, component: Component, part: bomFile.Part):
+        return apply_strategies(component, part, self.default_strategies)
+
+
+def work(sch, out_file, strategy: MakeBomStrategy, new_mode, pretty):
+    if not new_mode:
+        bom = to_bom_xml(sch)
+        xml = ElementTree.tostring(bom, encoding="unicode")
+
+        if pretty:
+            xml = minidom.parseString(xml).toprettyxml(indent="  ")
+
+        print(xml, file=out_file)
+    else:
+        file = bomFile.BomFile()
+
+        parts = bomFile.PartList()
+        file.partsProp = parts
+
+        components = to_bom(sch)
+        for c in components:
+            part = bomFile.Part(id=c.ref)
+            part.schema_reference = c.ref
+            part.part_numbersProp = bomFile.PartNumberList()
+
+            part = strategy.process_part(c, part)
+
+            if len(part.part_numbersProp.get_part_number()) == 0:
+                part.part_numbersProp = None
+
+            if part is not None:
+                parts.add_part(part)
+
+        file.export(out_file, 0, name_="bom-file", namespacedef_="xmlns='http://purl.org/ee/bom-file'",
+                    pretty_print=pretty)
+
+
+def make_bom(sch_file: str, out_file: Optional[str], strategy_name: str, new_mode: bool, pretty: bool):
+    sch = read_schematics(sch_file)
+
+    import pydoc
+    make_bom_strategy_factory = pydoc.locate(strategy_name)
+
+    if not callable(make_bom_strategy_factory):
+        raise EeException("Not a callable: {}, is a {}".format(strategy_name, type(make_bom_strategy_factory)))
+
+    make_bom_strategy = make_bom_strategy_factory()  # type: MakeBomStrategy
+
+    if not isinstance(make_bom_strategy, MakeBomStrategy):
+        raise EeException("Not a MakeBomStrategy: {}, is a {}".format(strategy_name, type(make_bom_strategy)))
+
+    if out_file:
+        mk_parents(out_file)
+        with open(out_file, "w") as f:
+            work(sch, f, make_bom_strategy, new_mode, pretty)
+    else:
+        work(sch, sys.stdout, make_bom_strategy, new_mode, pretty)