1 files changed, 34 insertions, 27 deletions
diff --git a/src/ee/tools/digikey_download_facts.py b/src/ee/tools/digikey_download_facts.py
index f21f171..a1e242b 100644
--- a/src/ee/tools/digikey_download_facts.py
+++ b/src/ee/tools/digikey_download_facts.py
@@ -1,11 +1,12 @@
-from colors import color
 import argparse
-import sys
-import ee.digikey as dk
-import pandas
 from itertools import *
-import yaml
+
 import os.path
+import yaml
+from colors import color
+
+import ee.digikey as dk
+from ee.digikey import SearchResponseTypes, DigikeyProduct
 
 parser = argparse.ArgumentParser(description="Download facts about parts from Digi-Key")
 
@@ -31,29 +32,35 @@ args = parser.parse_args()
 digikey = dk.Digikey()
 client = dk.DigikeyClient(digikey, on_download=lambda s: print(color(s, 'grey')))
 
+
 def mpn_to_path(mpn):
-  return "{}/{}.yaml".format(args.out, mpn)
+    return "{}/{}.yaml".format(args.out, mpn)
+
+
+def on_product(p: DigikeyProduct):
+    y = p.to_yaml()
+    with open(mpn_to_path(p.mpn), "w") as f:
+        yaml.dump(y, f, encoding="utf-8", allow_unicode=True)
 
-def on_product(p):
-  y = p.to_yaml()
-  with open(mpn_to_path(p.mpn), "w") as f:
-    yaml.dump(y, f, encoding="utf-8", allow_unicode=True)
 
 for p in args.parts:
-  print(color("Searching for {}".format(p), "white"))
-  path = mpn_to_path(p)
-
-  if os.path.isfile(path) and not args.force:
-    continue
-    
-  response = client.search(p)
-
-  if not response:
-    print(color("Part not found", "orange"))
-  elif isinstance(response, dk.DigikeyProduct):
-    print(color("Found {}".format(response.mpn)))
-    on_product(response)
-  else:
-    for k, g in groupby(sorted(response.products), lambda p: p.mpn):
-      print(color("Found {}".format(k), "white"))
-      on_product(list(g)[0])
+    print(color("Searching for {}".format(p), "white"))
+    path = mpn_to_path(p)
+
+    if os.path.isfile(path) and not args.force:
+        continue
+
+    response = client.search(p)
+
+    if response.response_type == SearchResponseTypes.SINGLE:
+        p = response.products[0]
+        print(color("Found {}".format(p.mpn), "white"))
+        on_product(p)
+    elif response.response_type == SearchResponseTypes.MANY:
+        for k, g in groupby(sorted(response.products), lambda p: p.mpn):
+            print(color("Found {}".format(k), "white"))
+            on_product(list(g)[0])
+    elif response.response_type == SearchResponseTypes.TOO_MANY:
+        print(color("Too many results ({}), select a category first".format(response.count), 'red'))
+    elif response.response_type == SearchResponseTypes.NO_MATCHES:
+        print(color("Part not found", "orange"))