feat: adding type cast for send_command

Author: humbertogontijo

roborock/api.py

@@ -13,7 +13,7 @@
 import struct
 import time
 from random import randint
-from typing import Any, Callable, Coroutine, Optional, Type
+from typing import Any, Callable, Coroutine, Optional, Type, TypeVar
 
 import aiohttp
 
@@ -30,6 +30,7 @@
     ModelStatus,
     MultiMapsList,
     NetworkInfo,
+    RoborockBase,
     RoomMapping,
     S7MaxVStatus,
     SmartWashParams,
@@ -63,6 +64,7 @@
     RoborockCommand.GET_MAP_V1,
     RoborockCommand.GET_MULTI_MAP,
 ]
+RT = TypeVar("RT", bound=RoborockBase)
 
 
 def md5hex(message: str) -> str:
@@ -228,36 +230,32 @@ def _get_payload(
         )
         return request_id, timestamp, payload
 
-    async def send_command(self, method: RoborockCommand, params: Optional[list | dict] = None):
+    async def send_command(
+        self,
+        method: RoborockCommand,
+        params: Optional[list | dict] = None,
+        return_type: Optional[Type[RT]] = None,
+    ) -> RT:
         raise NotImplementedError
 
     @fallback_cache
     async def get_status(self) -> Status | None:
-        status = await self.send_command(RoborockCommand.GET_STATUS)
-        if isinstance(status, dict):
-            _cls: Type[Status] = ModelStatus.get(
-                self.device_info.model, S7MaxVStatus
-            )  # Default to S7 MAXV if we don't have the data
-            return _cls.from_dict(status)
-        return None
+        _cls: Type[Status] = ModelStatus.get(
+            self.device_info.model, S7MaxVStatus
+        )  # Default to S7 MAXV if we don't have the data
+        return await self.send_command(RoborockCommand.GET_STATUS, return_type=_cls)
 
     @fallback_cache
     async def get_dnd_timer(self) -> DnDTimer | None:
-        dnd_timer = await self.send_command(RoborockCommand.GET_DND_TIMER)
-        if isinstance(dnd_timer, dict):
-            return DnDTimer.from_dict(dnd_timer)
-        return None
+        return await self.send_command(RoborockCommand.GET_DND_TIMER, return_type=DnDTimer)
 
     @fallback_cache
     async def get_valley_electricity_timer(self) -> ValleyElectricityTimer | None:
-        valley_electricity_timer = await self.send_command(RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER)
-        if isinstance(valley_electricity_timer, dict):
-            return ValleyElectricityTimer.from_dict(valley_electricity_timer)
-        return None
+        return await self.send_command(RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER, return_type=ValleyElectricityTimer)
 
     @fallback_cache
     async def get_clean_summary(self) -> CleanSummary | None:
-        clean_summary = await self.send_command(RoborockCommand.GET_CLEAN_SUMMARY)
+        clean_summary: dict | list | int = await self.send_command(RoborockCommand.GET_CLEAN_SUMMARY)
         if isinstance(clean_summary, dict):
             return CleanSummary.from_dict(clean_summary)
         elif isinstance(clean_summary, list):
@@ -274,38 +272,23 @@ async def get_clean_summary(self) -> CleanSummary | None:
 
     @fallback_cache
     async def get_clean_record(self, record_id: int) -> CleanRecord | None:
-        clean_record = await self.send_command(RoborockCommand.GET_CLEAN_RECORD, [record_id])
-        if isinstance(clean_record, dict):
-            return CleanRecord.from_dict(clean_record)
-        return None
+        return await self.send_command(RoborockCommand.GET_CLEAN_RECORD, [record_id], return_type=CleanRecord)
 
     @fallback_cache
     async def get_consumable(self) -> Consumable | None:
-        consumable = await self.send_command(RoborockCommand.GET_CONSUMABLE)
-        if isinstance(consumable, dict):
-            return Consumable.from_dict(consumable)
-        return None
+        return await self.send_command(RoborockCommand.GET_CONSUMABLE, return_type=Consumable)
 
     @fallback_cache
     async def get_wash_towel_mode(self) -> WashTowelMode | None:
-        washing_mode = await self.send_command(RoborockCommand.GET_WASH_TOWEL_MODE)
-        if isinstance(washing_mode, dict):
-            return WashTowelMode.from_dict(washing_mode)
-        return None
+        return await self.send_command(RoborockCommand.GET_WASH_TOWEL_MODE, return_type=WashTowelMode)
 
     @fallback_cache
     async def get_dust_collection_mode(self) -> DustCollectionMode | None:
-        dust_collection = await self.send_command(RoborockCommand.GET_DUST_COLLECTION_MODE)
-        if isinstance(dust_collection, dict):
-            return DustCollectionMode.from_dict(dust_collection)
-        return None
+        return await self.send_command(RoborockCommand.GET_DUST_COLLECTION_MODE, return_type=DustCollectionMode)
 
     @fallback_cache
     async def get_smart_wash_params(self) -> SmartWashParams | None:
-        mop_wash_mode = await self.send_command(RoborockCommand.GET_SMART_WASH_PARAMS)
-        if isinstance(mop_wash_mode, dict):
-            return SmartWashParams.from_dict(mop_wash_mode)
-        return None
+        return await self.send_command(RoborockCommand.GET_SMART_WASH_PARAMS, return_type=SmartWashParams)
 
     @fallback_cache
     async def get_dock_summary(self, dock_type: RoborockDockTypeCode) -> DockSummary | None:
@@ -359,22 +342,16 @@ async def get_prop(self) -> DeviceProp | None:
 
     @fallback_cache
     async def get_multi_maps_list(self) -> MultiMapsList | None:
-        multi_maps_list = await self.send_command(RoborockCommand.GET_MULTI_MAPS_LIST)
-        if isinstance(multi_maps_list, dict):
-            return MultiMapsList.from_dict(multi_maps_list)
-        return None
+        return await self.send_command(RoborockCommand.GET_MULTI_MAPS_LIST, return_type=MultiMapsList)
 
     @fallback_cache
     async def get_networking(self) -> NetworkInfo | None:
-        networking_info = await self.send_command(RoborockCommand.GET_NETWORK_INFO)
-        if isinstance(networking_info, dict):
-            return NetworkInfo.from_dict(networking_info)
-        return None
+        return await self.send_command(RoborockCommand.GET_NETWORK_INFO, return_type=NetworkInfo)
 
     @fallback_cache
     async def get_room_mapping(self) -> list[RoomMapping] | None:
         """Gets the mapping from segment id -> iot id. Only works on local api."""
-        mapping = await self.send_command(RoborockCommand.GET_ROOM_MAPPING)
+        mapping: list = await self.send_command(RoborockCommand.GET_ROOM_MAPPING)
         if isinstance(mapping, list):
             return [
                 RoomMapping(segment_id=segment_id, iot_id=iot_id)  # type: ignore
@@ -385,10 +362,7 @@ async def get_room_mapping(self) -> list[RoomMapping] | None:
     @fallback_cache
     async def get_child_lock_status(self) -> ChildLockStatus | None:
         """Gets current child lock status."""
-        child_lock_status = await self.send_command(RoborockCommand.GET_CHILD_LOCK_STATUS)
-        if isinstance(child_lock_status, dict):
-            return ChildLockStatus.from_dict(child_lock_status)
-        return None
+        return await self.send_command(RoborockCommand.GET_CHILD_LOCK_STATUS, return_type=ChildLockStatus)
 
     @fallback_cache
     async def get_sound_volume(self) -> int | None:

roborock/cloud_api.py

@@ -5,12 +5,12 @@
 import threading
 import uuid
 from asyncio import Lock
-from typing import Optional
+from typing import Optional, Type
 from urllib.parse import urlparse
 
 import paho.mqtt.client as mqtt
 
-from .api import COMMANDS_SECURED, KEEPALIVE, RoborockClient, md5hex
+from .api import COMMANDS_SECURED, KEEPALIVE, RT, RoborockClient, md5hex
 from .containers import DeviceData, UserData
 from .exceptions import CommandVacuumError, RoborockException, VacuumError
 from .protocol import MessageParser, Utils
@@ -149,7 +149,12 @@ def _send_msg_raw(self, msg: bytes) -> None:
         if info.rc != mqtt.MQTT_ERR_SUCCESS:
             raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")
 
-    async def send_command(self, method: RoborockCommand, params: Optional[list | dict] = None):
+    async def send_command(
+        self,
+        method: RoborockCommand,
+        params: Optional[list | dict] = None,
+        return_type: Optional[Type[RT]] = None,
+    ):
         await self.validate_connection()
         request_id, timestamp, payload = super()._get_payload(method, params, True)
         _LOGGER.debug(f"id={request_id} Requesting method {method} with {params}")
@@ -166,11 +171,9 @@ async def send_command(self, method: RoborockCommand, params: Optional[list | di
             _LOGGER.debug(f"id={request_id} Response from {method}: {len(response)} bytes")
         else:
             _LOGGER.debug(f"id={request_id} Response from {method}: {response}")
+        if return_type:
+            return return_type.from_dict(response)
         return response
 
     async def get_map_v1(self):
-        try:
-            return await self.send_command(RoborockCommand.GET_MAP_V1)
-        except RoborockException as e:
-            _LOGGER.error(e)
-        return None
+        return await self.send_command(RoborockCommand.GET_MAP_V1)

roborock/containers.py

@@ -80,8 +80,9 @@ class RoborockBase:
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]):
-        ignore_keys = cls._ignore_keys
-        return from_dict(cls, decamelize_obj(data, ignore_keys), config=Config(cast=[Enum]))
+        if isinstance(data, dict):
+            ignore_keys = cls._ignore_keys
+            return from_dict(cls, decamelize_obj(data, ignore_keys), config=Config(cast=[Enum]))
 
     def as_dict(self) -> dict:
         return asdict(
@@ -348,7 +349,9 @@ class DnDTimer(RoborockBase):
 
     def __post_init__(self) -> None:
         self.start_time = (
-            time(hour=self.start_hour, minute=self.start_minute) if self.start_hour and self.start_minute else None
+            time(hour=self.start_hour, minute=self.start_minute)
+            if self.start_hour is not None and self.start_minute is not None
+            else None
         )
         self.end_time = (
             time(hour=self.end_hour, minute=self.end_minute)
@@ -369,7 +372,9 @@ class ValleyElectricityTimer(RoborockBase):
 
     def __post_init__(self) -> None:
         self.start_time = (
-            time(hour=self.start_hour, minute=self.start_minute) if self.start_hour and self.start_minute else None
+            time(hour=self.start_hour, minute=self.start_minute)
+            if self.start_hour is not None and self.start_minute is not None
+            else None
         )
         self.end_time = (
             time(hour=self.end_hour, minute=self.end_minute)

roborock/local_api.py

@@ -3,13 +3,17 @@
 import asyncio
 import logging
 from asyncio import Lock, TimerHandle, Transport
-from typing import Optional
+from typing import Optional, Type
 
 import async_timeout
 
 from . import DeviceData
-from .api import COMMANDS_SECURED, QUEUE_TIMEOUT, RoborockClient
-from .exceptions import CommandVacuumError, RoborockConnectionException, RoborockException
+from .api import COMMANDS_SECURED, QUEUE_TIMEOUT, RT, RoborockClient
+from .exceptions import (
+    CommandVacuumError,
+    RoborockConnectionException,
+    RoborockException,
+)
 from .protocol import MessageParser
 from .roborock_message import RoborockMessage, RoborockMessageProtocol
 from .roborock_typing import RoborockCommand
@@ -123,9 +127,17 @@ async def ping(self):
             )
         )
 
-    async def send_command(self, method: RoborockCommand, params: Optional[list | dict] = None):
+    async def send_command(
+        self,
+        method: RoborockCommand,
+        params: Optional[list | dict] = None,
+        return_type: Optional[Type[RT]] = None,
+    ):
         roborock_message = self.build_roborock_message(method, params)
-        return (await self.send_message(roborock_message))[0]
+        response = (await self.send_message(roborock_message))[0]
+        if return_type:
+            return return_type.from_dict(response)
+        return response
 
     async def async_local_response(self, roborock_message: RoborockMessage):
         method = roborock_message.get_method()

tests/test_api.py

@@ -12,7 +12,7 @@
 )
 from roborock.api import PreparedRequest, RoborockApiClient
 from roborock.cloud_api import RoborockMqttClient
-from roborock.containers import DeviceData, S7MaxVStatus
+from roborock.containers import DeviceData, DustCollectionMode, S7MaxVStatus, SmartWashParams, WashTowelMode
 from tests.mock_data import BASE_URL_REQUEST, GET_CODE_RESPONSE, HOME_DATA_RAW, STATUS, USER_DATA
 
 
@@ -78,7 +78,7 @@ async def test_get_dust_collection_mode():
     device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
     with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
-        command.return_value = {"mode": 1}
+        command.return_value = DustCollectionMode.from_dict({"mode": 1})
         dust = await rmc.get_dust_collection_mode()
         assert dust is not None
         assert dust.mode == RoborockDockDustCollectionModeCode.light
@@ -90,7 +90,7 @@ async def test_get_mop_wash_mode():
     device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
     with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
-        command.return_value = {"smart_wash": 0, "wash_interval": 1500}
+        command.return_value = SmartWashParams.from_dict({"smart_wash": 0, "wash_interval": 1500})
         mop_wash = await rmc.get_smart_wash_params()
         assert mop_wash is not None
         assert mop_wash.smart_wash == 0
@@ -103,7 +103,7 @@ async def test_get_washing_mode():
     device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
     with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
-        command.return_value = {"wash_mode": 2}
+        command.return_value = WashTowelMode.from_dict({"wash_mode": 2})
         washing_mode = await rmc.get_wash_towel_mode()
         assert washing_mode is not None
         assert washing_mode.wash_mode == RoborockDockWashTowelModeCode.deep