"""Implementation of vacuum clean module.""" from __future__ import annotations import logging from enum import IntEnum from typing import Annotated from ...feature import Feature from ...module import FeatureAttribute from ..smartmodule import SmartModule _LOGGER = logging.getLogger(__name__) class Status(IntEnum): """Status of vacuum.""" Idle = 0 Cleaning = 1 Mapping = 2 GoingHome = 4 Charging = 5 Charged = 6 Paused = 7 Undocked = 8 Error = 100 UnknownInternal = -1000 class ErrorCode(IntEnum): """Error codes for vacuum.""" Ok = 0 SideBrushStuck = 2 MainBrushStuck = 3 WheelBlocked = 4 DustBinRemoved = 14 UnableToMove = 15 LidarBlocked = 16 UnableToFindDock = 21 BatteryLow = 22 UnknownInternal = -1000 class FanSpeed(IntEnum): """Fan speed level.""" Quiet = 1 Standard = 2 Turbo = 3 Max = 4 class Clean(SmartModule): """Implementation of vacuum clean module.""" REQUIRED_COMPONENT = "clean" _error_code = ErrorCode.Ok def _initialize_features(self) -> None: """Initialize features.""" self._add_feature( Feature( self._device, id="vacuum_return_home", name="Return home", container=self, attribute_setter="return_home", category=Feature.Category.Primary, type=Feature.Action, ) ) self._add_feature( Feature( self._device, id="vacuum_start", name="Start cleaning", container=self, attribute_setter="start", category=Feature.Category.Primary, type=Feature.Action, ) ) self._add_feature( Feature( self._device, id="vacuum_pause", name="Pause", container=self, attribute_setter="pause", category=Feature.Category.Primary, type=Feature.Action, ) ) self._add_feature( Feature( self._device, id="vacuum_status", name="Vacuum status", container=self, attribute_getter="status", category=Feature.Category.Primary, type=Feature.Type.Sensor, ) ) self._add_feature( Feature( self._device, id="vacuum_error", name="Error", container=self, attribute_getter="error", category=Feature.Category.Info, type=Feature.Type.Sensor, ) ) self._add_feature( Feature( self._device, id="battery_level", name="Battery level", container=self, attribute_getter="battery", icon="mdi:battery", unit_getter=lambda: "%", category=Feature.Category.Info, type=Feature.Type.Sensor, ) ) self._add_feature( Feature( self._device, id="vacuum_fan_speed", name="Fan speed", container=self, attribute_getter="fan_speed_preset", attribute_setter="set_fan_speed_preset", icon="mdi:fan", choices_getter=lambda: list(FanSpeed.__members__), category=Feature.Category.Primary, type=Feature.Type.Choice, ) ) async def _post_update_hook(self) -> None: """Set error code after update.""" errors = self._vac_status.get("err_status") if errors is None or not errors: self._error_code = ErrorCode.Ok return if len(errors) > 1: _LOGGER.warning( "Multiple error codes, using the first one only: %s", errors ) error = errors.pop(0) try: self._error_code = ErrorCode(error) except ValueError: _LOGGER.warning( "Unknown error code, please create an issue describing the error: %s", error, ) self._error_code = ErrorCode.UnknownInternal def query(self) -> dict: """Query to execute during the update cycle.""" return { "getVacStatus": None, "getBatteryInfo": None, "getCleanStatus": None, "getCleanAttr": {"type": "global"}, } async def start(self) -> dict: """Start cleaning.""" # If we are paused, do not restart cleaning if self.status is Status.Paused: return await self.resume() return await self.call( "setSwitchClean", { "clean_mode": 0, "clean_on": True, "clean_order": True, "force_clean": False, }, ) async def pause(self) -> dict: """Pause cleaning.""" if self.status is Status.GoingHome: return await self.set_return_home(False) return await self.set_pause(True) async def resume(self) -> dict: """Resume cleaning.""" return await self.set_pause(False) async def set_pause(self, enabled: bool) -> dict: """Pause or resume cleaning.""" return await self.call("setRobotPause", {"pause": enabled}) async def return_home(self) -> dict: """Return home.""" return await self.set_return_home(True) async def set_return_home(self, enabled: bool) -> dict: """Return home / pause returning.""" return await self.call("setSwitchCharge", {"switch_charge": enabled}) @property def error(self) -> ErrorCode: """Return error.""" return self._error_code @property def fan_speed_preset(self) -> Annotated[str, FeatureAttribute()]: """Return fan speed preset.""" return FanSpeed(self._settings["suction"]).name async def set_fan_speed_preset( self, speed: str ) -> Annotated[dict, FeatureAttribute]: """Set fan speed preset.""" name_to_value = {x.name: x.value for x in FanSpeed} if speed not in name_to_value: raise ValueError("Invalid fan speed %s, available %s", speed, name_to_value) return await self.call( "setCleanAttr", {"suction": name_to_value[speed], "type": "global"} ) @property def battery(self) -> int: """Return battery level.""" return self.data["getBatteryInfo"]["battery_percentage"] @property def _vac_status(self) -> dict: """Return vac status container.""" return self.data["getVacStatus"] @property def _settings(self) -> dict: """Return cleaning settings.""" return self.data["getCleanAttr"] @property def status(self) -> Status: """Return current status.""" if self._error_code is not ErrorCode.Ok: return Status.Error status_code = self._vac_status["status"] try: return Status(status_code) except ValueError: _LOGGER.warning("Got unknown status code: %s (%s)", status_code, self.data) return Status.UnknownInternal