roborock.devices.traits.a01

  1"""Create traits for A01 devices.
  2
  3This module provides the API implementations for A01 protocol devices, which include
  4Dyad (Wet/Dry Vacuums) and Zeo (Washing Machines).
  5
  6Using A01 APIs
  7--------------
  8A01 devices expose a single API object that handles all device interactions. This API is
  9available on the device instance (typically via `device.a01_properties`).
 10
 11The API provides two main methods:
 121.  **query_values(protocols)**: Fetches current state for specific data points.
 13    You must pass a list of protocol enums (e.g. `RoborockDyadDataProtocol` or
 14    `RoborockZeoProtocol`) to request specific data.
 152.  **set_value(protocol, value)**: Sends a command to the device to change a setting
 16    or perform an action.
 17
 18Note that these APIs fetch data directly from the device upon request and do not
 19cache state internally.
 20"""
 21
 22import json
 23from collections.abc import Callable
 24from datetime import time
 25from typing import Any
 26
 27from roborock.data import DyadProductInfo, DyadSndState, HomeDataProduct, RoborockCategory
 28from roborock.data.dyad.dyad_code_mappings import (
 29    DyadBrushSpeed,
 30    DyadCleanMode,
 31    DyadError,
 32    DyadSelfCleanLevel,
 33    DyadSelfCleanMode,
 34    DyadSuction,
 35    DyadWarmLevel,
 36    DyadWaterLevel,
 37    RoborockDyadStateCode,
 38)
 39from roborock.data.zeo.zeo_code_mappings import (
 40    ZeoDetergentType,
 41    ZeoDryingMode,
 42    ZeoError,
 43    ZeoMode,
 44    ZeoProgram,
 45    ZeoRinse,
 46    ZeoSoftenerType,
 47    ZeoSpin,
 48    ZeoState,
 49    ZeoTemperature,
 50)
 51from roborock.devices.a01_channel import send_decoded_command
 52from roborock.devices.mqtt_channel import MqttChannel
 53from roborock.devices.traits import Trait
 54from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
 55
 56__init__ = [
 57    "DyadApi",
 58    "ZeoApi",
 59]
 60
 61
 62DYAD_PROTOCOL_ENTRIES: dict[RoborockDyadDataProtocol, Callable] = {
 63    RoborockDyadDataProtocol.STATUS: lambda val: RoborockDyadStateCode(val).name,
 64    RoborockDyadDataProtocol.SELF_CLEAN_MODE: lambda val: DyadSelfCleanMode(val).name,
 65    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: lambda val: DyadSelfCleanLevel(val).name,
 66    RoborockDyadDataProtocol.WARM_LEVEL: lambda val: DyadWarmLevel(val).name,
 67    RoborockDyadDataProtocol.CLEAN_MODE: lambda val: DyadCleanMode(val).name,
 68    RoborockDyadDataProtocol.SUCTION: lambda val: DyadSuction(val).name,
 69    RoborockDyadDataProtocol.WATER_LEVEL: lambda val: DyadWaterLevel(val).name,
 70    RoborockDyadDataProtocol.BRUSH_SPEED: lambda val: DyadBrushSpeed(val).name,
 71    RoborockDyadDataProtocol.POWER: lambda val: int(val),
 72    RoborockDyadDataProtocol.AUTO_DRY: lambda val: bool(val),
 73    RoborockDyadDataProtocol.MESH_LEFT: lambda val: int(360000 - val * 60),
 74    RoborockDyadDataProtocol.BRUSH_LEFT: lambda val: int(360000 - val * 60),
 75    RoborockDyadDataProtocol.ERROR: lambda val: DyadError(val).name,
 76    RoborockDyadDataProtocol.VOLUME_SET: lambda val: int(val),
 77    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: lambda val: bool(val),
 78    RoborockDyadDataProtocol.AUTO_DRY_MODE: lambda val: bool(val),
 79    RoborockDyadDataProtocol.SILENT_DRY_DURATION: lambda val: int(val),  # in minutes
 80    RoborockDyadDataProtocol.SILENT_MODE: lambda val: bool(val),
 81    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: lambda val: time(
 82        hour=int(val / 60), minute=val % 60
 83    ),  # in minutes since 00:00
 84    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: lambda val: time(
 85        hour=int(val / 60), minute=val % 60
 86    ),  # in minutes since 00:00
 87    RoborockDyadDataProtocol.RECENT_RUN_TIME: lambda val: [
 88        int(v) for v in val.split(",")
 89    ],  # minutes of cleaning in past few days.
 90    RoborockDyadDataProtocol.TOTAL_RUN_TIME: lambda val: int(val),
 91    RoborockDyadDataProtocol.SND_STATE: lambda val: DyadSndState.from_dict(val),
 92    RoborockDyadDataProtocol.PRODUCT_INFO: lambda val: DyadProductInfo.from_dict(val),
 93}
 94
 95ZEO_PROTOCOL_ENTRIES: dict[RoborockZeoProtocol, Callable] = {
 96    # read-only
 97    RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name,
 98    RoborockZeoProtocol.COUNTDOWN: lambda val: int(val),
 99    RoborockZeoProtocol.WASHING_LEFT: lambda val: int(val),
100    RoborockZeoProtocol.ERROR: lambda val: ZeoError(val).name,
101    RoborockZeoProtocol.TIMES_AFTER_CLEAN: lambda val: int(val),
102    RoborockZeoProtocol.DETERGENT_EMPTY: lambda val: bool(val),
103    RoborockZeoProtocol.SOFTENER_EMPTY: lambda val: bool(val),
104    # read-write
105    RoborockZeoProtocol.MODE: lambda val: ZeoMode(val).name,
106    RoborockZeoProtocol.PROGRAM: lambda val: ZeoProgram(val).name,
107    RoborockZeoProtocol.TEMP: lambda val: ZeoTemperature(val).name,
108    RoborockZeoProtocol.RINSE_TIMES: lambda val: ZeoRinse(val).name,
109    RoborockZeoProtocol.SPIN_LEVEL: lambda val: ZeoSpin(val).name,
110    RoborockZeoProtocol.DRYING_MODE: lambda val: ZeoDryingMode(val).name,
111    RoborockZeoProtocol.DETERGENT_TYPE: lambda val: ZeoDetergentType(val).name,
112    RoborockZeoProtocol.SOFTENER_TYPE: lambda val: ZeoSoftenerType(val).name,
113    RoborockZeoProtocol.SOUND_SET: lambda val: bool(val),
114}
115
116
117def convert_dyad_value(protocol_value: RoborockDyadDataProtocol, value: Any) -> Any:
118    """Convert a dyad protocol value to its corresponding type."""
119    if (converter := DYAD_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
120        try:
121            return converter(value)
122        except (ValueError, TypeError):
123            return None
124    return None
125
126
127def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any:
128    """Convert a zeo protocol value to its corresponding type."""
129    if (converter := ZEO_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
130        try:
131            return converter(value)
132        except (ValueError, TypeError):
133            return None
134    return None
135
136
137class DyadApi(Trait):
138    """API for interacting with Dyad devices."""
139
140    def __init__(self, channel: MqttChannel) -> None:
141        """Initialize the Dyad API."""
142        self._channel = channel
143
144    async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[RoborockDyadDataProtocol, Any]:
145        """Query the device for the values of the given Dyad protocols."""
146        response = await send_decoded_command(
147            self._channel,
148            {RoborockDyadDataProtocol.ID_QUERY: protocols},
149            value_encoder=json.dumps,
150        )
151        return {protocol: convert_dyad_value(protocol, response.get(protocol)) for protocol in protocols}
152
153    async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dict[RoborockDyadDataProtocol, Any]:
154        """Set a value for a specific protocol on the device."""
155        params = {protocol: value}
156        return await send_decoded_command(self._channel, params)
157
158
159class ZeoApi(Trait):
160    """API for interacting with Zeo devices."""
161
162    name = "zeo"
163
164    def __init__(self, channel: MqttChannel) -> None:
165        """Initialize the Zeo API."""
166        self._channel = channel
167
168    async def query_values(self, protocols: list[RoborockZeoProtocol]) -> dict[RoborockZeoProtocol, Any]:
169        """Query the device for the values of the given protocols."""
170        response = await send_decoded_command(
171            self._channel,
172            {RoborockZeoProtocol.ID_QUERY: protocols},
173            value_encoder=json.dumps,
174        )
175        return {protocol: convert_zeo_value(protocol, response.get(protocol)) for protocol in protocols}
176
177    async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[RoborockZeoProtocol, Any]:
178        """Set a value for a specific protocol on the device."""
179        params = {protocol: value}
180        return await send_decoded_command(self._channel, params, value_encoder=lambda x: x)
181
182
183def create(product: HomeDataProduct, mqtt_channel: MqttChannel) -> DyadApi | ZeoApi:
184    """Create traits for A01 devices."""
185    match product.category:
186        case RoborockCategory.WET_DRY_VAC:
187            return DyadApi(mqtt_channel)
188        case RoborockCategory.WASHING_MACHINE:
189            return ZeoApi(mqtt_channel)
190        case _:
191            raise NotImplementedError(f"Unsupported category {product.category}")