Adjust callback type

homewizard · Nov 17, 2024 · 3c2a743 · 3c2a743
1 parent e9443e7
commit 3c2a743
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/homewizard_energy/v2/websocket.py b/homewizard_energy/v2/websocket.py
@@ -14,6 +14,8 @@
 if TYPE_CHECKING:
     from . import HomeWizardEnergyV2
 
+OnMessageCallbackType = Callable[[str, Device | Measurement | System], None]
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -22,7 +24,8 @@ class Websocket:
 
     _connect_lock: asyncio.Lock = asyncio.Lock()
     _ws_connection: aiohttp.ClientWebSocketResponse | None = None
-    _ws_subscriptions: list[tuple[str, Callable[[aiohttp.WSMessage], None]]] = []
+    _ws_subscriptions: list[tuple[str, OnMessageCallbackType]] = []
+
     _ws_authenticated: bool = False
 
     def __init__(self, parent: "HomeWizardEnergyV2"):
@@ -65,7 +68,7 @@ async def disconnect(self) -> None:
         self._ws_connection = None
 
     def subscribe(
-        self, topic: WebsocketTopic, ws_callback: Callable[[aiohttp.WSMessage], None]
+        self, topic: WebsocketTopic, ws_callback: OnMessageCallbackType
     ) -> Callable[[], None]:
         """
         Subscribe to raw websocket messages.