1 files changed, 94 insertions, 0 deletions
diff --git a/src/ee/digikey/search_parts.py b/src/ee/digikey/search_parts.py
new file mode 100644
index 0000000..62bb4ee
--- /dev/null
+++ b/src/ee/digikey/search_parts.py
@@ -0,0 +1,94 @@
+from pathlib import Path
+from typing import List
+
+from ee.digikey import Digikey, DigikeyParser, DigikeyClient, SearchResponseTypes, DigikeyProduct
+from ee.part import PartDb, load_db, save_db
+from ee.xml import bomFile, bom_file_utils
+from ee.xml.bomFile import DigikeyDistributorInfo
+from ee.xml.uris import DIGIKEY_URI
+
+__all__ = ["search_parts"]
+
+
+def resolved(di: DigikeyDistributorInfo, part: bomFile.Part, p: DigikeyProduct):
+    di.stateProp = "resolved"
+
+    fact_set = bom_file_utils.find_fact_set(part, DIGIKEY_URI, create=True)
+
+    # Remove the old list
+    fact_set.factsProp = bomFile.FactList()
+    facts: List[bomFile.Fact] = fact_set.factsProp.factProp
+
+    for a in p.attributes:
+        facts.append(bomFile.Fact(key=a.attribute_type.id, label=a.attribute_type.label, value=a.value))
+
+
+def search_parts(in_dir: Path, out_dir: Path, cache_dir: Path, force_refresh: bool):
+    print("in: {}, out: {}".format(in_dir, out_dir))
+
+    in_db = load_db(in_dir)
+    out_parts = PartDb()
+
+    parser = DigikeyParser(Digikey())
+    client = DigikeyClient(cache_dir)
+
+    for part in in_db.iterparts():
+        dpn = bom_file_utils.find_dpn(part, DIGIKEY_URI)
+        mpn = bom_file_utils.find_pn(part)
+
+        is_mpn = query = None
+
+        if dpn is not None:
+            query = dpn
+            is_mpn = False
+        elif mpn is not None:
+            query = mpn
+            is_mpn = True
+
+        if query is None:
+            print("could not find pn or dpn: part.id={}".format(part.idProp))
+            continue
+
+        di = part.distributor_infoProp  # type: DigikeyDistributorInfo
+
+        if di is None:
+            di = bomFile.DigikeyDistributorInfo()
+            di.extensiontype_ = "DigikeyDistributorInfo"
+            di.original_tagname_ = "distributor-info"
+            part.distributor_infoProp = di
+
+        if force_refresh or di.stateProp != "resolved":
+            text = client.search(query)
+            response = parser.parse_string(text)
+
+            if response.response_type == SearchResponseTypes.SINGLE:
+                resolved(di, part, response.products[0])
+            elif response.response_type == SearchResponseTypes.MANY:
+
+                # find those with an exact match. Digikey uses a prefix search so a query for "FOO" will return "FOO"
+                # and "FOOT".
+                def get_field(p):
+                    return p.mpn if is_mpn else p.part_number
+
+                filtered_products = [p for p in response.products if get_field(p) == query]
+
+                if len(filtered_products) == 0:
+                    di.stateProp = "not-found"
+                else:
+                    dpn = sorted(filtered_products, key=lambda p: p.part_number)[0].part_number
+
+                    response = parser.parse_string(client.search(dpn))
+                    if response.response_type == SearchResponseTypes.SINGLE:
+                        resolved(di, part, response.products[0])
+                    else:
+                        di.stateProp = "many"
+
+            elif response.response_type == SearchResponseTypes.TOO_MANY:
+                di.stateProp = "too-many"
+            elif response.response_type == SearchResponseTypes.NO_MATCHES:
+                di.stateProp = "not-found"
+
+        out_parts.add_entry(part, True)
+
+    print("Saving {} work parts".format(out_parts.size()))
+    save_db(out_dir, out_parts)