Skip to content

Commit

Permalink
fix: mypy complaints (#137)
Browse files Browse the repository at this point in the history
* fix: change some typing

* fix: include poetry lock

* fix: linting

* fix: linting

* fix: add typing

* fix: linting

* fix: bugs

* fix: linting

* fix: none typing

* Update api.py
  • Loading branch information
Lash-L authored Nov 3, 2023
1 parent d2d53a0 commit 752e320
Show file tree
Hide file tree
Showing 6 changed files with 65 additions and 20 deletions.
32 changes: 22 additions & 10 deletions roborock/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,7 @@ def _async_response(
def _get_payload(
self,
method: RoborockCommand | str,
params: list | dict | None = None,
params: list | dict | int | None = None,
secured=False,
):
timestamp = math.floor(time.time())
Expand Down Expand Up @@ -381,15 +381,15 @@ async def send_message(self, roborock_message: RoborockMessage):
async def _send_command(
self,
method: RoborockCommand | str,
params: list | dict | None = None,
params: list | dict | int | None = None,
):
raise NotImplementedError

@final
async def send_command(
self,
method: RoborockCommand | str,
params: list | dict | None = None,
params: list | dict | int | None = None,
return_type: type[RT] | None = None,
) -> RT:
cacheable_attribute_result = find_cacheable_attribute(method)
Expand All @@ -412,8 +412,11 @@ async def send_command(
return return_type.from_dict(response)
return response

async def get_status(self) -> Status | None:
return self._status_type.from_dict(await self.cache[CacheableAttribute.status].async_value())
async def get_status(self) -> Status:
data = self._status_type.from_dict(await self.cache[CacheableAttribute.status].async_value())
if data is None:
return self._status_type()
return data

async def get_dnd_timer(self) -> DnDTimer | None:
return DnDTimer.from_dict(await self.cache[CacheableAttribute.dnd_timer].async_value())
Expand Down Expand Up @@ -451,8 +454,11 @@ async def get_clean_record(self, record_id: int) -> CleanRecord | None:
_LOGGER.warning("Clean record was of a new type, please submit an issue request: %s", record)
return None

async def get_consumable(self) -> Consumable | None:
return Consumable.from_dict(await self.cache[CacheableAttribute.consumable].async_value())
async def get_consumable(self) -> Consumable:
data = Consumable.from_dict(await self.cache[CacheableAttribute.consumable].async_value())
if data is None:
return Consumable()
return data

async def get_wash_towel_mode(self) -> WashTowelMode | None:
return WashTowelMode.from_dict(await self.cache[CacheableAttribute.wash_towel_mode].async_value())
Expand All @@ -463,7 +469,7 @@ async def get_dust_collection_mode(self) -> DustCollectionMode | None:
async def get_smart_wash_params(self) -> SmartWashParams | None:
return SmartWashParams.from_dict(await self.cache[CacheableAttribute.smart_wash_params].async_value())

async def get_dock_summary(self, dock_type: RoborockDockTypeCode) -> DockSummary | None:
async def get_dock_summary(self, dock_type: RoborockDockTypeCode) -> DockSummary:
"""Gets the status summary from the dock with the methods available for a given dock.
:param dock_type: RoborockDockTypeCode"""
Expand Down Expand Up @@ -525,11 +531,11 @@ async def get_room_mapping(self) -> list[RoomMapping] | None:
]
return None

async def get_child_lock_status(self) -> ChildLockStatus | None:
async def get_child_lock_status(self) -> ChildLockStatus:
"""Gets current child lock status."""
return ChildLockStatus.from_dict(await self.cache[CacheableAttribute.child_lock_status].async_value())

async def get_flow_led_status(self) -> FlowLedStatus | None:
async def get_flow_led_status(self) -> FlowLedStatus:
"""Gets current flow led status."""
return FlowLedStatus.from_dict(await self.cache[CacheableAttribute.flow_led_status].async_value())

Expand All @@ -549,6 +555,12 @@ async def get_server_timer(self) -> list[ServerTimer]:
def add_listener(self, listener: Callable):
self._listeners.append(listener)

async def get_from_cache(self, key: CacheableAttribute) -> AttributeCache | None:
val = self.cache.get(key)
if val is not None:
return await val.async_value()
return None


class RoborockApiClient:
def __init__(self, username: str, base_url=None) -> None:
Expand Down
2 changes: 1 addition & 1 deletion roborock/cloud_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,7 @@ async def send_message(self, roborock_message: RoborockMessage):
async def _send_command(
self,
method: RoborockCommand | str,
params: list | dict | None = None,
params: list | dict | int | None = None,
):
request_id, timestamp, payload = super()._get_payload(method, params, True)
request_protocol = RoborockMessageProtocol.RPC_REQUEST
Expand Down
2 changes: 1 addition & 1 deletion roborock/const.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
ROBOROCK_G10S_PRO = "roborock.vacuum.a26"
ROBOROCK_G10S = "roborock.vacuum.a46"
ROBOROCK_G10 = "roborock.vacuum.a29"
ROCKROBO_G10_SG = "roborock.vacuum.a30" # Variant of the G10, has similar features as S7
ROCKROBO_G10_SG = "roborock.vacuum.a30" # Variant of the G10, has similar features as S7
ROBOROCK_S7 = "roborock.vacuum.a15"
ROBOROCK_S6_MAXV = "roborock.vacuum.a10"
ROBOROCK_E2 = "roborock.vacuum.e2"
Expand Down
35 changes: 34 additions & 1 deletion roborock/containers.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
SENSOR_DIRTY_REPLACE_TIME,
SIDE_BRUSH_REPLACE_TIME,
)
from .exceptions import RoborockException

_LOGGER = logging.getLogger(__name__)

Expand Down Expand Up @@ -185,12 +186,12 @@ class HomeDataDevice(RoborockBase):
name: str
local_key: str
fv: str
product_id: str
attribute: Any | None = None
active_time: int | None = None
runtime_env: Any | None = None
time_zone_id: str | None = None
icon_url: str | None = None
product_id: str | None = None
lon: Any | None = None
lat: Any | None = None
share: Any | None = None
Expand Down Expand Up @@ -297,9 +298,41 @@ class Status(RoborockBase):
dss: int | None = None
common_status: int | None = None
corner_clean_mode: int | None = None
error_code_name: str | None = None
state_name: str | None = None
water_box_mode_name: str | None = None
fan_power_options: list[str] = field(default_factory=list)
fan_power_name: str | None = None
mop_mode_name: str | None = None

def __post_init__(self) -> None:
self.square_meter_clean_area = round(self.clean_area / 1000000, 1) if self.clean_area is not None else None
if self.error_code is not None:
self.error_code_name = self.error_code.name
if self.state is not None:
self.state_name = self.state.name
if self.water_box_mode is not None:
self.water_box_mode_name = self.water_box_mode.name
if self.fan_power is not None:
self.fan_power_options = self.fan_power.keys()
self.fan_power_name = self.fan_power.name
if self.mop_mode is not None:
self.mop_mode_name = self.mop_mode.name

def get_fan_speed_code(self, fan_speed: str) -> int:
if self.fan_power is None:
raise RoborockException("Attempted to get fan speed before status has been updated.")
return self.fan_power.as_dict().get(fan_speed)

def get_mop_intensity_code(self, mop_intensity: str) -> int:
if self.water_box_mode is None:
raise RoborockException("Attempted to get mop_intensity before status has been updated.")
return self.water_box_mode.as_dict().get(mop_intensity)

def get_mop_mode_code(self, mop_mode: str) -> int:
if self.mop_mode is None:
raise RoborockException("Attempted to get mop_mode before status has been updated.")
return self.mop_mode.as_dict().get(mop_mode)


@dataclass
Expand Down
4 changes: 2 additions & 2 deletions roborock/local_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ async def async_disconnect(self) -> None:
self.sync_disconnect()

def build_roborock_message(
self, method: RoborockCommand | str, params: list | dict | None = None
self, method: RoborockCommand | str, params: list | dict | int | None = None
) -> RoborockMessage:
secured = True if method in COMMANDS_SECURED else False
request_id, timestamp, payload = self._get_payload(method, params, secured)
Expand Down Expand Up @@ -123,7 +123,7 @@ async def ping(self):
async def _send_command(
self,
method: RoborockCommand | str,
params: list | dict | None = None,
params: list | dict | int | None = None,
):
roborock_message = self.build_roborock_message(method, params)
return await self.send_message(roborock_message)
Expand Down
10 changes: 5 additions & 5 deletions roborock/roborock_typing.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from dataclasses import dataclass
from dataclasses import dataclass, field
from enum import Enum

from .containers import (
Expand Down Expand Up @@ -132,7 +132,7 @@ class RoborockCommand(str, Enum):

@dataclass
class CommandInfo:
params: list | dict | None = None
params: list | dict | int | None = None


CommandInfoMap: dict[RoborockCommand | None, CommandInfo] = {
Expand Down Expand Up @@ -315,9 +315,9 @@ class DockSummary(RoborockBase):

@dataclass
class DeviceProp(RoborockBase):
status: Status | None = None
clean_summary: CleanSummary | None = None
consumable: Consumable | None = None
status: Status = field(default_factory=Status)
clean_summary: CleanSummary = field(default_factory=CleanSummary)
consumable: Consumable = field(default_factory=Consumable)
last_clean_record: CleanRecord | None = None
dock_summary: DockSummary | None = None

Expand Down

0 comments on commit 752e320

Please sign in to comment.