feat: improve some typing (#204)

humbertogontijo · May 9, 2024 · 7752db9 · 7752db9
1 parent 459782e
commit 7752db9
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 5 deletions.
diff --git a/roborock/api.py b/roborock/api.py
@@ -49,10 +49,10 @@ def __init__(self, endpoint: str, device_info: DeviceData, queue_timeout: int =
     def __del__(self) -> None:
         self.release()
 
-    def release(self):
+    def release(self) -> None:
         self.sync_disconnect()
 
-    async def async_release(self):
+    async def async_release(self) -> None:
         await self.async_disconnect()
 
     @property

diff --git a/roborock/version_1_apis/roborock_mqtt_client_v1.py b/roborock/version_1_apis/roborock_mqtt_client_v1.py
@@ -71,5 +71,5 @@ async def _send_command(
         roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload)
         return await self.send_message(roborock_message)
 
-    async def get_map_v1(self):
+    async def get_map_v1(self) -> bytes | None:
         return await self.send_command(RoborockCommand.GET_MAP_V1)
diff --git a/roborock/version_a01_apis/roborock_client_a01.py b/roborock/version_a01_apis/roborock_client_a01.py
@@ -137,6 +137,8 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                         if queue and queue.protocol == protocol:
                             queue.resolve((converted_response, None))
 
-    async def update_values(self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]):
+    async def update_values(
+        self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]
+    ) -> dict[RoborockDyadDataProtocol | RoborockZeoProtocol, typing.Any]:
         """This should handle updating for each given protocol."""
         raise NotImplementedError
diff --git a/roborock/version_a01_apis/roborock_mqtt_client_a01.py b/roborock/version_a01_apis/roborock_mqtt_client_a01.py
@@ -57,7 +57,9 @@ async def send_message(self, roborock_message: RoborockMessage):
                     dps_responses[dps] = response[0]
         return dps_responses
 
-    async def update_values(self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]):
+    async def update_values(
+        self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]
+    ) -> dict[RoborockDyadDataProtocol | RoborockZeoProtocol, typing.Any]:
         payload = {"dps": {RoborockDyadDataProtocol.ID_QUERY: str([int(protocol) for protocol in dyad_data_protocols])}}
         return await self.send_message(
             RoborockMessage(