fix: typing for generic BleakClient classes and the retry_bluetooth_c…

…onnection_error decorator (#86)

* fix: typing for the generic BleakClient client class

Using a bound TypeVar we can ensure that any client class we are dealing
with is either BleakClient or a descendant of it and that type then
stays consistent throughout the lifecycle.

Signed-off-by: Felix Kaechele <felix@kaechele.ca>

* fix: typing for the retry_bluetooth_connection_error decorator

Use TypeVar together with ParamVar to drop the use of the unsafe cast
operation.

Signed-off-by: Felix Kaechele <felix@kaechele.ca>

---------

Signed-off-by: Felix Kaechele <felix@kaechele.ca>
Co-authored-by: J. Nick Koston <nick@koston.org>

src/bleak_retry_connector/__init__.py

@@ -1,14 +1,11 @@
 from __future__ import annotations
 
-from typing import cast
-
 __version__ = "2.13.1"
 
 
 import asyncio
 import logging
-from collections.abc import Callable
-from typing import Any, TypeVar
+from typing import Any, Awaitable, Callable, ParamSpec, TypeVar
 
 import async_timeout
 from bleak import BleakClient, BleakScanner
@@ -279,17 +276,20 @@ async def close_stale_connections(
     await _disconnect_devices(to_disconnect)
 
 
+AnyBleakClient = TypeVar("AnyBleakClient", bound=BleakClient)
+
+
 async def establish_connection(
-    client_class: type[BleakClient],
+    client_class: type[AnyBleakClient],
     device: BLEDevice,
     name: str,
-    disconnected_callback: Callable[[BleakClient], None] | None = None,
+    disconnected_callback: Callable[[AnyBleakClient], None] | None = None,
     max_attempts: int = MAX_CONNECT_ATTEMPTS,
     cached_services: BleakGATTServiceCollection | None = None,
     ble_device_callback: Callable[[], BLEDevice] | None = None,
     use_services_cache: bool = True,
     **kwargs: Any,
-) -> BleakClient:
+) -> AnyBleakClient:
     """Establish a connection to the device."""
     timeouts = 0
     connect_errors = 0
@@ -427,22 +427,27 @@ def _raise_if_needed(name: str, description: str, exc: Exception) -> None:
     raise RuntimeError("This should never happen")
 
 
-WrapFuncType = TypeVar("WrapFuncType", bound=Callable[..., Any])
+P = ParamSpec("P")
+T = TypeVar("T")
 
 
-def retry_bluetooth_connection_error(attempts: int = DEFAULT_ATTEMPTS) -> WrapFuncType:  # type: ignore[type-var]
+def retry_bluetooth_connection_error(
+    attempts: int = DEFAULT_ATTEMPTS,
+) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
     """Define a wrapper to retry on bluetooth connection error."""
 
-    def _decorator_retry_bluetooth_connection_error(func: WrapFuncType) -> WrapFuncType:
+    def _decorator_retry_bluetooth_connection_error(
+        func: Callable[P, Awaitable[T]]
+    ) -> Callable[P, Awaitable[T]]:
         """Define a wrapper to retry on bleak error.
 
         The accessory is allowed to disconnect us any time so
         we need to retry the operation.
         """
 
-        async def _async_wrap_bluetooth_connection_error_retry(
-            *args: Any, **kwargs: Any
-        ) -> Any:
+        async def _async_wrap_bluetooth_connection_error_retry(  # type: ignore[return]
+            *args: P.args, **kwargs: P.kwargs
+        ) -> T:
             for attempt in range(attempts):
                 try:
                     return await func(*args, **kwargs)
@@ -458,9 +463,9 @@ async def _async_wrap_bluetooth_connection_error_retry(
                     )
                     await asyncio.sleep(backoff_time)
 
-        return cast(WrapFuncType, _async_wrap_bluetooth_connection_error_retry)
+        return _async_wrap_bluetooth_connection_error_retry
 
-    return cast(WrapFuncType, _decorator_retry_bluetooth_connection_error)
+    return _decorator_retry_bluetooth_connection_error
 
 
 async def restore_discoveries(scanner: BleakScanner, adapter: str) -> None:

tests/test_init.py

@@ -1427,7 +1427,7 @@ def test_calculate_backoff_time():
 async def test_retry_bluetooth_connection_error():
     """Test that the retry_bluetooth_connection_error decorator works correctly."""
 
-    @retry_bluetooth_connection_error()  # type: ignore[misc]
+    @retry_bluetooth_connection_error()
     async def test_function():
         raise BleakDBusError(MagicMock(), MagicMock())
 
@@ -1445,7 +1445,7 @@ async def test_function():
 async def test_retry_bluetooth_connection_error_non_default_max_attempts():
     """Test that the retry_bluetooth_connection_error decorator works correctly with a different number of retries."""
 
-    @retry_bluetooth_connection_error(4)  # type: ignore[misc]
+    @retry_bluetooth_connection_error(4)
     async def test_function():
         raise BleakDBusError(MagicMock(), MagicMock())