From 1a3f4dab365ca5478c4494f032046a45a53d674c Mon Sep 17 00:00:00 2001
From: MattHag <16444067+MattHag@users.noreply.github.com>
Date: Sat, 28 Dec 2024 23:51:55 +0100
Subject: [PATCH] Speedup lookup of known receivers

Refactor get_receiver_info. Replacing data structure of known receivers
to avoid for loop, when an efficient dictionary lookup is possible.

Related #2273
---
 lib/logitech_receiver/base_usb.py        | 85 ++++++++++++++----------
 tests/logitech_receiver/test_base_usb.py | 30 +++++++++
 2 files changed, 79 insertions(+), 36 deletions(-)
 create mode 100644 tests/logitech_receiver/test_base_usb.py

diff --git a/lib/logitech_receiver/base_usb.py b/lib/logitech_receiver/base_usb.py
index 41e572b8..be960dd7 100644
--- a/lib/logitech_receiver/base_usb.py
+++ b/lib/logitech_receiver/base_usb.py
@@ -27,6 +27,10 @@ USB ids of Logitech wireless receivers.
 Only receivers supporting the HID++ protocol can go in here.
 """
 
+from __future__ import annotations
+
+from typing import Any
+
 from solaar.i18n import _
 
 # max_devices is only used for receivers that do not support reading from Registers.RECEIVER_INFO offset 0x03, default
@@ -174,49 +178,58 @@ LIGHTSPEED_RECEIVER_C547 = _lightspeed_receiver(0xC547)
 # EX100 old style receiver pre-unifying protocol
 EX100_27MHZ_RECEIVER_C517 = _ex100_receiver(0xC517)
 
-KNOWN_RECEIVERS = (
-    BOLT_RECEIVER_C548,
-    UNIFYING_RECEIVER_C52B,
-    UNIFYING_RECEIVER_C532,
-    NANO_RECEIVER_ADVANCED,
-    NANO_RECEIVER_C518,
-    NANO_RECEIVER_C51A,
-    NANO_RECEIVER_C51B,
-    NANO_RECEIVER_C521,
-    NANO_RECEIVER_C525,
-    NANO_RECEIVER_C526,
-    NANO_RECEIVER_C52E,
-    NANO_RECEIVER_C531,
-    NANO_RECEIVER_C534,
-    NANO_RECEIVER_C535,
-    NANO_RECEIVER_C537,
-    NANO_RECEIVER_6042,
-    LIGHTSPEED_RECEIVER_C539,
-    LIGHTSPEED_RECEIVER_C53A,
-    LIGHTSPEED_RECEIVER_C53D,
-    LIGHTSPEED_RECEIVER_C53F,
-    LIGHTSPEED_RECEIVER_C541,
-    LIGHTSPEED_RECEIVER_C545,
-    LIGHTSPEED_RECEIVER_C547,
-    EX100_27MHZ_RECEIVER_C517,
-)
+KNOWN_RECEIVERS = {
+    0xC548: BOLT_RECEIVER_C548,
+    0xC52B: UNIFYING_RECEIVER_C52B,
+    0xC532: UNIFYING_RECEIVER_C532,
+    0xC52F: NANO_RECEIVER_ADVANCED,
+    0xC518: NANO_RECEIVER_C518,
+    0xC51A: NANO_RECEIVER_C51A,
+    0xC51B: NANO_RECEIVER_C51B,
+    0xC521: NANO_RECEIVER_C521,
+    0xC525: NANO_RECEIVER_C525,
+    0xC526: NANO_RECEIVER_C526,
+    0xC52E: NANO_RECEIVER_C52E,
+    0xC531: NANO_RECEIVER_C531,
+    0xC534: NANO_RECEIVER_C534,
+    0xC535: NANO_RECEIVER_C535,
+    0xC537: NANO_RECEIVER_C537,
+    0x6042: NANO_RECEIVER_6042,
+    0xC539: LIGHTSPEED_RECEIVER_C539,
+    0xC53A: LIGHTSPEED_RECEIVER_C53A,
+    0xC53D: LIGHTSPEED_RECEIVER_C53D,
+    0xC53F: LIGHTSPEED_RECEIVER_C53F,
+    0xC541: LIGHTSPEED_RECEIVER_C541,
+    0xC545: LIGHTSPEED_RECEIVER_C545,
+    0xC547: LIGHTSPEED_RECEIVER_C547,
+    0xC517: EX100_27MHZ_RECEIVER_C517,
+}
 
 
-def get_receiver_info(product_id: int) -> dict:
-    """Returns hardcoded information about Logitech receiver.
+def get_receiver_info(product_id: int) -> dict[str, Any]:
+    """Returns hardcoded information about a Logitech receiver.
 
     Parameters
     ----------
     product_id
-        Product ID of receiver e.g. 0xC548 for a Logitech Bolt receiver.
+        Product ID (pid) of the receiver, e.g. 0xC548 for a Logitech
+        Bolt receiver.
 
     Returns
     -------
-    dict
-        Product info with mandatory vendor_id, product_id,
-        usb_interface, name, receiver_kind
+    dict[str, Any]
+        Receiver info with mandatory fields:
+        - vendor_id
+        - product_id
+
+    Raises
+    ------
+    ValueError
+        If the product ID is unknown.
     """
-    for receiver in KNOWN_RECEIVERS:
-        if product_id == receiver.get("product_id"):
-            return receiver
-    raise ValueError(f"Unknown product ID '0x{product_id:02X}")
+    try:
+        return KNOWN_RECEIVERS[product_id]
+    except KeyError:
+        pass
+
+    raise ValueError(f"Unknown product ID '0x{product_id:02X}'")
diff --git a/tests/logitech_receiver/test_base_usb.py b/tests/logitech_receiver/test_base_usb.py
new file mode 100644
index 00000000..b911f7ed
--- /dev/null
+++ b/tests/logitech_receiver/test_base_usb.py
@@ -0,0 +1,30 @@
+import pytest
+
+from logitech_receiver import base_usb
+from logitech_receiver.common import LOGITECH_VENDOR_ID
+
+
+def test_ensure_known_receivers_mappings_are_valid():
+    for key, receiver in base_usb.KNOWN_RECEIVERS.items():
+        assert key == receiver["product_id"]
+
+
+def test_get_receiver_info():
+    expected = {
+        "vendor_id": LOGITECH_VENDOR_ID,
+        "product_id": 0xC548,
+        "usb_interface": 2,
+        "name": "Bolt Receiver",
+        "receiver_kind": "bolt",
+        "max_devices": 6,
+        "may_unpair": True,
+    }
+
+    res = base_usb.get_receiver_info(0xC548)
+
+    assert res == expected
+
+
+def test_get_receiver_info_unknown_device_fails():
+    with pytest.raises(ValueError):
+        base_usb.get_receiver_info(0xC500)