api: Update type checking with use of generic Device type

ZacheryThomas · ZacheryThomas · commit e7fd0d9c01bb · 2024-12-03T15:05:54.000-05:00
diff --git a/smartrent/api.py b/smartrent/api.py
@@ -1,8 +1,9 @@
 import logging
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, TypeVar
 
 import aiohttp
 
+from .device import Device
 from .lock import DoorLock
 from .sensor import LeakSensor
 from .switch import BinarySwitch, MultilevelSwitch
@@ -12,9 +13,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    ALL_DEVICE_TYPES = Union[
-        DoorLock, Thermostat, BinarySwitch, MultilevelSwitch, LeakSensor
-    ]
+    ALL_DEVICE_TYPES = TypeVar("ALL_DEVICE_TYPES", bound=Device)
 
 
 class API:
@@ -37,7 +36,7 @@ def __init__(
         aiohttp_session: aiohttp.ClientSession = None,
         tfa_token=None,
     ):
-        self._device_list: List["ALL_DEVICE_TYPES"] = []
+        self._device_list: List["ALL_DEVICE_TYPES"] = []  # type: ignore
         self.client: Client = Client(email, password, aiohttp_session, tfa_token)
 
     async def async_fetch_devices(self):
