Add datasheet support for all suppliers (#1)

inventree_part_import/config/__init__.py

@@ -105,12 +105,19 @@ def setup_inventree_api():
     r"^(?P<scheme>[^:/\s]+://)?(?P<hostname>[^:/\s]+)(?::(?P<port>\d{1,5}))?(?P<path>/.*)?$")
 
 DEFAULT_CONFIG_VARS = {
-    "max_results": 10,
     "interactive": "twice",
+    "max_results": 10,
     "request_timeout": 15.0,
     "retry_timeout": 3.0,
 }
-VALID_CONFIG_VARS = {"currency", "language", "location", "scraping", *DEFAULT_CONFIG_VARS}
+VALID_CONFIG_VARS = {
+    "currency",
+    "language",
+    "location",
+    "scraping",
+    "datasheets",
+    *DEFAULT_CONFIG_VARS,
+}
 
 _CONFIG_LOADED = None
 CONFIG = "config.yaml"
@@ -143,16 +150,26 @@ def get_config(reload=False):
     currency = input_currency()
     language = input_language()
     location = input_location()
+
     prompt("do you want to enable web scraping? (this is required to use some suppliers)",
         prefix="", end="\n")
     warning("enabling scraping can get you temporarily blocked sometimes")
     scraping = prompt_yes_or_no("enable scraping?", default_is_yes=True)
 
+    prompt("how do you want to handle datasheets?")
+    datasheets_choices = [
+        "upload (upload file attachments for parts)",
+        "false  (do not add datasheets for parts)",
+    ]
+    datasheets_values = ["upload", False]
+    datasheets_index = select(datasheets_choices, deselected_prefix="  ", selected_prefix="> ")
+
     _CONFIG_LOADED = {
         "currency": currency,
         "language": language,
         "location": location,
         "scraping": scraping,
+        "datasheets": datasheets_values[datasheets_index],
         **DEFAULT_CONFIG_VARS,
     }
     yaml_data = yaml_dump(_CONFIG_LOADED, sort_keys=False)

inventree_part_import/inventree_helpers.py

@@ -1,5 +1,7 @@
+from base64 import urlsafe_b64encode
 from dataclasses import dataclass
 from functools import cache
+from hashlib import sha256
 import re
 
 from inventree.api import InvenTreeAPI
@@ -9,7 +11,7 @@
 from inventree.part import ParameterTemplate, Part
 from platformdirs import user_cache_path
 import requests
-from requests.compat import urlparse
+from requests.compat import unquote, urlparse
 from requests.exceptions import HTTPError, Timeout
 
 from .error_helper import *
@@ -93,28 +95,49 @@ def download_image_content(api_object: ImageMixin):
 
 def upload_image(api_object: ImageMixin, image_url: str):
     info("uploading image ...")
-    if not (image_content := _download_image_content(image_url)):
+    image_content, redirected_url = _download_file_content(image_url)
+    if not image_content:
         warning(f"failed to download image from '{image_url}'")
         return
 
-    url_path = urlparse(image_url).path
-    if "." in url_path:
-        file_extension = url_path.rsplit(".")[-1]
-    else:
-        file_extension = image_url.rsplit(".")[-1]
+    file_extension = url2filename(redirected_url).split(".")[-1]
     if not file_extension.isalnum():
         warning(f"failed to get file extension for image from '{image_url}'")
         return
 
-    image_path = INVENTREE_CACHE / f"temp_image.{file_extension}"
-    with open(image_path, "wb") as file:
-        file.write(image_content)
+    image_hash = urlsafe_b64encode(sha256(image_content).digest()).decode()
+    image_path = INVENTREE_CACHE / f"{image_hash}.{file_extension}"
+    image_path.write_bytes(image_content)
 
     try:
         api_object.uploadImage(str(image_path))
     except HTTPError as e:
         warning(f"failed to upload image with: {e.args[0]['body']}")
 
+def upload_datasheet(part: Part, datasheet_url: str):
+    info("uploading datasheet ...")
+    datasheet_content, redirected_url = _download_file_content(datasheet_url)
+    if not datasheet_content:
+        warning(f"failed to download datasheet from '{datasheet_url}'")
+        return
+
+    file_name = url2filename(redirected_url)
+    file_extension = file_name.split(".")[-1]
+    if file_extension != "pdf":
+        warning(f"datasheet '{datasheet_url}' has invalid file extension '{file_extension}'")
+        return
+
+    datasheet_path = INVENTREE_CACHE / file_name
+    datasheet_path.write_bytes(datasheet_content)
+
+    try:
+        part.uploadAttachment(str(datasheet_path), "datasheet")
+    except HTTPError as e:
+        warning(f"failed to upload datasheet with: {e.args[0]['body']}")
+
+def url2filename(url):
+    return unquote(urlparse(url).path.split("/")[-1])
+
 DOWNLOAD_HEADERS = {"User-Agent": "Mozilla/5.0"}
 
 from ssl import PROTOCOL_TLSv1_2
@@ -129,7 +152,7 @@ def init_poolmanager(self, connections, maxsize, block=False):
         )
 
 @cache
-def _download_image_content(url):
+def _download_file_content(url):
     session = requests.Session()
     session.mount("https://", TLSv1_2HTTPAdapter())
 
@@ -139,10 +162,10 @@ def _download_image_content(url):
                 result = session.get(url, headers=DOWNLOAD_HEADERS)
                 result.raise_for_status()
     except (HTTPError, Timeout) as e:
-        warning(f"failed to download image with '{e}'")
+        warning(f"failed to download file with '{e}'")
         return None
 
-    return result.content
+    return result.content, result.url
 
 @dataclass
 class Company:

inventree_part_import/part_importer.py

@@ -9,11 +9,11 @@
 from thefuzz import fuzz
 
 from .categories import setup_categories_and_parameters
-from .config import CATEGORIES_CONFIG, get_config, get_pre_creation_hooks
+from .config import CATEGORIES_CONFIG, CONFIG, get_config, get_pre_creation_hooks
 from .error_helper import *
 from .inventree_helpers import (create_manufacturer, get_manufacturer_part,
                                 get_parameter_templates, get_part, get_supplier_part,
-                                update_object_data, upload_image)
+                                update_object_data, upload_datasheet, upload_image)
 from .suppliers import search
 from .suppliers.base import ApiPart
 
@@ -144,6 +144,16 @@ def import_supplier_part(self, supplier: Company, api_part: ApiPart):
             if not part.image and api_part.image_url:
                 upload_image(part, api_part.image_url)
 
+            attachment_types = {attachment.comment for attachment in part.getAttachments()}
+            if "datasheet" not in attachment_types and api_part.datasheet_url:
+                match get_config().get("datasheets"):
+                    case "upload":
+                        upload_datasheet(part, api_part.datasheet_url)
+                    case None | False:
+                        pass
+                    case invalid_mode:
+                        warning(f"invalid value 'datasheets: {invalid_mode}' in {CONFIG}")
+
         if api_part.parameters:
             result = self.setup_parameters(part, api_part, update_part)
             import_result |= result

inventree_part_import/suppliers/base.py

@@ -10,6 +10,7 @@
 class ApiPart:
     description: str
     image_url: str
+    datasheet_url: str
     supplier_link: str
     SKU: str
     manufacturer: str

inventree_part_import/suppliers/supplier_digikey.py

@@ -103,6 +103,7 @@ def get_api_part(self, digikey_part):
         return ApiPart(
             description=digikey_part.product_description,
             image_url=digikey_part.primary_photo,
+            datasheet_url=digikey_part.primary_datasheet,
             supplier_link=digikey_part.product_url,
             SKU=digikey_part.digi_key_part_number,
             manufacturer=digikey_part.manufacturer.value,

inventree_part_import/suppliers/supplier_lcsc.py

@@ -116,6 +116,7 @@ def get_api_part(self, lcsc_part):
         return ApiPart(
             description=REMOVE_HTML_TAGS.sub("", description),
             image_url=image_url,
+            datasheet_url=lcsc_part.get("pdfUrl"),
             supplier_link=supplier_link,
             SKU=lcsc_part.get("productCode", ""),
             manufacturer=lcsc_part.get("brandNameEn", ""),

inventree_part_import/suppliers/supplier_mouser.py

@@ -74,6 +74,7 @@ def get_api_part(self, mouser_part):
         api_part = ApiPart(
             description=REMOVE_HTML_TAGS.sub("", mouser_part.get("Description", "")),
             image_url=mouser_part.get("ImagePath"),
+            datasheet_url=mouser_part.get("DataSheetUrl"),
             supplier_link=supplier_link,
             SKU=mouser_part_number,
             manufacturer=mouser_part.get("Manufacturer", ""),

inventree_part_import/suppliers/supplier_reichelt.py

@@ -3,7 +3,7 @@
 
 from bs4 import BeautifulSoup
 from requests import Session
-from requests.compat import quote
+from requests.compat import quote, urljoin
 
 from ..config import get_config
 from ..error_helper import *
@@ -79,7 +79,12 @@ def search(self, search_term):
 
     def get_api_part(self, soup, sku, link):
         description = soup.find(id="av_articleheader").find("span", itemprop="name").text
-        img_url = soup.find(id="av_bildbox").find(id="bigimages nohighlight").find("img")["src"]
+
+        bigimage = soup.find(id="av_bildbox").find(id="bigimages nohighlight")
+        image_url = bigimage.find("img")["src"] if bigimage else None
+
+        datasheet = soup.find(id="av_datasheetview").find(class_="av_datasheet")
+        datasheet_url = urljoin(BASE_URL, datasheet.find("a")["href"]) if datasheet else None
 
         availability = soup.find("p", class_="availability").find("span")["class"][0]
         if availability not in AVAILABILITY_MAP:
@@ -120,7 +125,8 @@ def get_api_part(self, soup, sku, link):
 
         return ApiPart(
             description=description,
-            image_url=img_url,
+            image_url=image_url,
+            datasheet_url=datasheet_url,
             supplier_link=link,
             SKU=sku.upper(),
             manufacturer=manufacturer,

inventree_part_import/suppliers/supplier_tme.py

@@ -52,6 +52,7 @@ def get_api_part(self, tme_part, tme_stock):
         api_part = ApiPart(
             description=tme_part.get("Description", ""),
             image_url=fix_tme_url(tme_part.get("Photo")),
+            datasheet_url=None,
             supplier_link=fix_tme_url(tme_part.get("ProductInformationPage")),
             SKU=tme_part.get("Symbol", ""),
             manufacturer=tme_part.get("Producer", ""),
@@ -72,12 +73,17 @@ def get_api_part(self, tme_part, tme_stock):
     def finalize_hook(self, api_part: ApiPart):
         if not (parameters := self.tme_api.get_parameters(api_part.SKU)):
             return False
-
         api_part.parameters = {
             parameter["ParameterName"]: parameter["ParameterValue"]
             for parameter in parameters
         }
 
+        if product_files := self.tme_api.get_product_files(api_part.SKU):
+            for document in product_files.get("DocumentList", []):
+                if document.get("DocumentType") == "DTE":
+                    api_part.datasheet_url = fix_tme_url(document.get("DocumentUrl"))
+                    break
+
         return True
 
 def fix_tme_url(url):
@@ -159,6 +165,16 @@ def get_prices_and_stocks(self, product_symbols):
             return result_data["ProductList"]
         return []
 
+    def get_categories(self):
+        result = self._api_call("Products/GetCategories", {
+            "Country": self.country,
+            "Language": self.language,
+            "Tree": "false",
+        })
+        if result:
+            return result.json()["Data"]["CategoryTree"]
+        return []
+
     def get_parameters(self, product_symbol):
         result = self._api_call("Products/GetParameters", {
             "Country": self.country,
@@ -169,14 +185,14 @@ def get_parameters(self, product_symbol):
             return result.json()["Data"]["ProductList"][0]["ParameterList"]
         return []
 
-    def get_categories(self):
-        result = self._api_call("Products/GetCategories", {
+    def get_product_files(self, product_symbol):
+        result = self._api_call("Products/GetProductsFiles", {
             "Country": self.country,
             "Language": self.language,
-            "Tree": "false",
+            "SymbolList[0]": product_symbol,
         })
         if result:
-            return result.json()["Data"]["CategoryTree"]
+            return result.json()["Data"]["ProductList"][0]["Files"]
         return []
 
     HEADERS = {"Content-type": "application/x-www-form-urlencoded"}

tests/test_cli.py

@@ -39,7 +39,7 @@ def setup_class(self):
         self.api = InvenTreeAPI(HOST, username=USERNAME, password=PASSWORD, use_token_auth=True)
 
         (TEST_CONFIG_DIR / CONFIG).write_text(
-            "currency: EUR\nlanguage: EN\nlocation: DE\nscraping: true\n")
+            "currency: EUR\nlanguage: EN\nlocation: DE\nscraping: true\ndatasheets: upload\n")
         (TEST_CONFIG_DIR / INVENTREE_CONFIG).write_text(
             f"host: {HOST}\ntoken: {self.api.token}\n")
         (TEST_CONFIG_DIR / SUPPLIERS_CONFIG).write_text(