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 (`device.dyad` or `device.zeo`).
 10
 11The API provides these 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.
 173.  **values**: The latest known state, merged from query responses and unsolicited
 18    pushes in arrival order.
 194.  **add_update_listener(callback)**: Registers a callback invoked whenever `values`
 20    changes; read `values` from the callback to get the updated state.
 21
 22The device pushes only the data points that changed, so `values` is the merged view
 23of everything seen so far. State tracking is active once the device is connected
 24(the device calls `start()` on the API, which subscribes to the MQTT topic).
 25"""
 26
 27import json
 28import logging
 29from abc import abstractmethod
 30from collections.abc import Callable
 31from datetime import UTC, datetime, time
 32from typing import Any, Generic, TypeVar
 33
 34from roborock.data import DyadProductInfo, DyadSndState, HomeDataProduct, RoborockCategory
 35from roborock.data.dyad.dyad_code_mappings import (
 36    DyadBrushSpeed,
 37    DyadCleanMode,
 38    DyadError,
 39    DyadSelfCleanLevel,
 40    DyadSelfCleanMode,
 41    DyadSuction,
 42    DyadWarmLevel,
 43    DyadWaterLevel,
 44    RoborockDyadStateCode,
 45)
 46from roborock.data.zeo.zeo_code_mappings import (
 47    ZeoDetergentType,
 48    ZeoDryingMode,
 49    ZeoError,
 50    ZeoFeatureBits,
 51    ZeoMode,
 52    ZeoProgram,
 53    ZeoRinse,
 54    ZeoSoftenerType,
 55    ZeoSpin,
 56    ZeoState,
 57    ZeoTemperature,
 58)
 59from roborock.devices.rpc.a01_channel import send_decoded_command
 60from roborock.devices.traits import Trait
 61from roborock.devices.traits.a01.device_feature import (
 62    build_feature_dp_list,
 63    build_force_load_dp_list,
 64    supports_uv_light,
 65)
 66from roborock.devices.traits.common import TraitUpdateListener
 67from roborock.devices.transport.mqtt_channel import MqttChannel
 68from roborock.exceptions import RoborockException
 69from roborock.protocols.a01_protocol import decode_rpc_response
 70from roborock.roborock_message import (
 71    RoborockDyadDataProtocol,
 72    RoborockMessage,
 73    RoborockMessageProtocol,
 74    RoborockZeoProtocol,
 75)
 76
 77_LOGGER = logging.getLogger(__name__)
 78
 79__all__ = [
 80    "A01Api",
 81    "DyadApi",
 82    "ZeoApi",
 83]
 84
 85
 86DYAD_PROTOCOL_ENTRIES: dict[RoborockDyadDataProtocol, Callable] = {
 87    RoborockDyadDataProtocol.STATUS: lambda val: RoborockDyadStateCode(val).name,
 88    RoborockDyadDataProtocol.SELF_CLEAN_MODE: lambda val: DyadSelfCleanMode(val).name,
 89    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: lambda val: DyadSelfCleanLevel(val).name,
 90    RoborockDyadDataProtocol.WARM_LEVEL: lambda val: DyadWarmLevel(val).name,
 91    RoborockDyadDataProtocol.CLEAN_MODE: lambda val: DyadCleanMode(val).name,
 92    RoborockDyadDataProtocol.SUCTION: lambda val: DyadSuction(val).name,
 93    RoborockDyadDataProtocol.WATER_LEVEL: lambda val: DyadWaterLevel(val).name,
 94    RoborockDyadDataProtocol.BRUSH_SPEED: lambda val: DyadBrushSpeed(val).name,
 95    RoborockDyadDataProtocol.POWER: lambda val: int(val),
 96    RoborockDyadDataProtocol.AUTO_DRY: lambda val: bool(val),
 97    RoborockDyadDataProtocol.MESH_LEFT: lambda val: int(360000 - val * 60),
 98    RoborockDyadDataProtocol.BRUSH_LEFT: lambda val: int(360000 - val * 60),
 99    RoborockDyadDataProtocol.ERROR: lambda val: DyadError(val).name,
100    RoborockDyadDataProtocol.VOLUME_SET: lambda val: int(val),
101    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: lambda val: bool(val),
102    RoborockDyadDataProtocol.AUTO_DRY_MODE: lambda val: bool(val),
103    RoborockDyadDataProtocol.SILENT_DRY_DURATION: lambda val: int(val),  # in minutes
104    RoborockDyadDataProtocol.SILENT_MODE: lambda val: bool(val),
105    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: lambda val: time(
106        hour=int(val / 60), minute=val % 60
107    ),  # in minutes since 00:00
108    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: lambda val: time(
109        hour=int(val / 60), minute=val % 60
110    ),  # in minutes since 00:00
111    RoborockDyadDataProtocol.RECENT_RUN_TIME: lambda val: [
112        int(v) for v in val.split(",")
113    ],  # minutes of cleaning in past few days.
114    RoborockDyadDataProtocol.TOTAL_RUN_TIME: lambda val: int(val),
115    RoborockDyadDataProtocol.SND_STATE: lambda val: DyadSndState.from_dict(val),
116    RoborockDyadDataProtocol.PRODUCT_INFO: lambda val: DyadProductInfo.from_dict(val),
117}
118
119ZEO_PROTOCOL_ENTRIES: dict[RoborockZeoProtocol, Callable] = {
120    # read-only
121    RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name,
122    RoborockZeoProtocol.COUNTDOWN: lambda val: int(val),
123    RoborockZeoProtocol.WASHING_LEFT: lambda val: int(val),
124    RoborockZeoProtocol.ERROR: lambda val: ZeoError(val).name,
125    RoborockZeoProtocol.TIMES_AFTER_CLEAN: lambda val: int(val),
126    RoborockZeoProtocol.DETERGENT_EMPTY: lambda val: bool(val),
127    RoborockZeoProtocol.SOFTENER_EMPTY: lambda val: bool(val),
128    # read-write
129    RoborockZeoProtocol.MODE: lambda val: ZeoMode(val).name,
130    RoborockZeoProtocol.PROGRAM: lambda val: ZeoProgram(val).name,
131    RoborockZeoProtocol.TEMP: lambda val: ZeoTemperature(val).name,
132    RoborockZeoProtocol.RINSE_TIMES: lambda val: ZeoRinse(val).name,
133    RoborockZeoProtocol.SPIN_LEVEL: lambda val: ZeoSpin(val).name,
134    RoborockZeoProtocol.DRYING_MODE: lambda val: ZeoDryingMode(val).name,
135    RoborockZeoProtocol.DETERGENT_TYPE: lambda val: ZeoDetergentType(val).name,
136    RoborockZeoProtocol.SOFTENER_TYPE: lambda val: ZeoSoftenerType(val).name,
137    RoborockZeoProtocol.SOUND_SET: lambda val: bool(val),
138    RoborockZeoProtocol.FEATURE_BITS: lambda val: int(val),
139}
140
141
142def convert_dyad_value(protocol_value: RoborockDyadDataProtocol, value: Any) -> Any:
143    """Convert a dyad protocol value to its corresponding type."""
144    if (converter := DYAD_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
145        try:
146            return converter(value)
147        except (ValueError, TypeError):
148            return None
149    return None
150
151
152def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any:
153    """Convert a zeo protocol value to its corresponding type."""
154    if (converter := ZEO_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
155        try:
156            return converter(value)
157        except (ValueError, TypeError):
158            return None
159    return None
160
161
162_DYAD_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockDyadDataProtocol)
163_ZEO_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockZeoProtocol)
164
165_P = TypeVar("_P", RoborockDyadDataProtocol, RoborockZeoProtocol)
166
167
168class A01Api(Trait, TraitUpdateListener, Generic[_P]):
169    """Base class for A01 device APIs with device state tracking.
170
171    Query responses and unsolicited pushes both arrive on the same MQTT topic,
172    so a single subscription merges every decoded message into `values` in
173    arrival order. Update listeners are notified whenever a value changes.
174    """
175
176    def __init__(self, channel: MqttChannel, initial_status: dict[int, Any] | None = None) -> None:
177        """Initialize the A01 API, optionally seeding `values` from a cloud status snapshot."""
178        TraitUpdateListener.__init__(self, _LOGGER)
179        self._channel = channel
180        self._values: dict[_P, Any] = {}
181        self._unsub: Callable[[], None] | None = None
182        self._last_message_time: datetime | None = None
183        if initial_status:
184            self._merge_values(self._decode_datapoints(initial_status))
185
186    @property
187    def values(self) -> dict[_P, Any]:
188        """Latest known device state, merged from query responses and pushes.
189
190        The device pushes only the data points that changed, so this is the
191        merged view of everything seen so far. A protocol the device has not
192        reported yet is absent from the dictionary.
193        """
194        return dict(self._values)
195
196    @property
197    def last_message_time(self) -> datetime | None:
198        """Time the last message was received from the device.
199
200        Updated on every decoded message, even when no value changed: idle
201        devices push an identical heartbeat, so this is the liveness signal
202        even when `values` stays the same and update listeners stay silent.
203        The initial cloud status snapshot does not count as a message.
204        """
205        return self._last_message_time
206
207    async def start(self) -> None:
208        """Subscribe to the device state topic and start tracking `values`."""
209        await self._ensure_subscribed()
210
211    def close(self) -> None:
212        """Unsubscribe from MQTT push and release resources."""
213        if self._unsub is not None:
214            self._unsub()
215            self._unsub = None
216
217    async def _ensure_subscribed(self) -> None:
218        """Subscribe to MQTT DPS push (idempotent)."""
219        if self._unsub is not None:
220            return
221        self._unsub = await self._channel.subscribe(self._on_message)
222
223    @abstractmethod
224    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[_P, Any]:
225        """Convert raw datapoints to typed values, skipping unknown codes."""
226
227    def _on_message(self, message: RoborockMessage) -> None:
228        """Handle a message on the device topic (query response or push)."""
229        if message.protocol != RoborockMessageProtocol.RPC_RESPONSE:
230            return
231        try:
232            datapoints = decode_rpc_response(message)
233        except RoborockException:
234            _LOGGER.debug("Dropped malformed push message", exc_info=True)
235            return
236        self._last_message_time = datetime.now(UTC)
237        self._merge_values(self._decode_datapoints(datapoints))
238
239    def _merge_query_response(self, values: dict[_P, Any]) -> None:
240        """Record a successful query response when there is no subscription.
241
242        When subscribed, the response was already merged in arrival order and
243        timestamped by `_on_message`; merging again here could overwrite a
244        push that arrived after it.
245        """
246        if self._unsub is not None:
247            return
248        self._last_message_time = datetime.now(UTC)
249        self._merge_values(values)
250
251    def _merge_values(self, values: dict[_P, Any]) -> None:
252        """Merge decoded values into the cache and notify on change."""
253        changed = False
254        for protocol, value in values.items():
255            if value is None:
256                continue
257            if protocol not in self._values or self._values[protocol] != value:
258                self._values[protocol] = value
259                changed = True
260        if changed:
261            self._notify_update()
262
263
264class DyadApi(A01Api[RoborockDyadDataProtocol]):
265    """API for interacting with Dyad devices."""
266
267    name = "dyad"
268
269    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[RoborockDyadDataProtocol, Any]:
270        """Convert raw datapoints to typed values, skipping unknown codes."""
271        values: dict[RoborockDyadDataProtocol, Any] = {}
272        for code, value in datapoints.items():
273            if code not in _DYAD_PROTOCOL_VALUES:
274                continue
275            protocol = RoborockDyadDataProtocol(code)
276            values[protocol] = convert_dyad_value(protocol, value)
277        return values
278
279    async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[RoborockDyadDataProtocol, Any]:
280        """Query the device for the values of the given Dyad protocols."""
281        response = await send_decoded_command(
282            self._channel,
283            {RoborockDyadDataProtocol.ID_QUERY: protocols},
284            value_encoder=json.dumps,
285        )
286        values = {protocol: convert_dyad_value(protocol, response.get(protocol)) for protocol in protocols}
287        self._merge_query_response(values)
288        return values
289
290    async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dict[RoborockDyadDataProtocol, Any]:
291        """Set a value for a specific protocol on the device."""
292        params = {protocol: value}
293        return await send_decoded_command(self._channel, params)
294
295    async def add_listener(self, callback: Callable[[dict[RoborockDyadDataProtocol, Any]], None]) -> Callable[[], None]:
296        """Listen for state the device pushes on its own.
297
298        The callback is invoked with decoded values whenever the device sends a
299        message, including unsolicited pushes when its state changes. Only known
300        protocols are delivered. Returns a callable to remove the listener.
301
302        Prefer `add_update_listener` together with `values`, which handle the
303        merging of partial pushes for you.
304        """
305
306        def on_message(message: RoborockMessage) -> None:
307            try:
308                datapoints = decode_rpc_response(message)
309            except RoborockException:
310                return
311            if values := self._decode_datapoints(datapoints):
312                callback(values)
313
314        return await self._channel.subscribe(on_message)
315
316
317class ZeoApi(A01Api[RoborockZeoProtocol]):
318    """API for interacting with Zeo devices."""
319
320    name = "zeo"
321
322    def __init__(
323        self, channel: MqttChannel, model: str | None = None, initial_status: dict[int, Any] | None = None
324    ) -> None:
325        """Initialize the Zeo API."""
326        self._feature_bits: int = 0
327        self._model = model
328        super().__init__(channel, initial_status)
329
330    async def start(self) -> None:
331        """Subscribe to MQTT push and trigger a full state sync.
332
333        Subscribes to the DPS MQTT topic, then performs a two-stage
334        force-load: first the base DP list (including FEATURE_BITS),
335        then a second query for the DPs gated behind each enabled feature.
336        The device responds with a complete state dump;
337        subsequent changes arrive via incremental MQTT push.
338        """
339        await self._ensure_subscribed()
340        await self._force_load()
341        await self._load_feature_dps()
342
343    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[RoborockZeoProtocol, Any]:
344        """Convert raw datapoints to typed values, skipping unknown codes."""
345        values: dict[RoborockZeoProtocol, Any] = {}
346        for code, value in datapoints.items():
347            if code not in _ZEO_PROTOCOL_VALUES:
348                continue
349            protocol = RoborockZeoProtocol(code)
350            values[protocol] = convert_zeo_value(protocol, value)
351        return values
352
353    async def _force_load(self) -> None:
354        """Send ID_QUERY with the base DP list to trigger a full state push.
355
356        For devices known to lack FEATURE_BITS, the DP is excluded
357        from the query list and ``_feature_bits`` stays at 0.
358        """
359        dp_list = build_force_load_dp_list(self._model)
360        result = await self.query_values(dp_list)
361        self._feature_bits = result.get(RoborockZeoProtocol.FEATURE_BITS, 0)
362
363    async def _load_feature_dps(self) -> None:
364        """Second-stage query for feature-gated DPs.
365
366        Called unconditionally after the first force-load; each DP is
367        independently gated:
368
369        - Feature-gated DPs are queried only when their feature bit is set
370          in FEATURE_BITS (DP 237).
371        - UV light (DP 228) is gated by :func:`supports_uv_light` (series
372          whitelist), independent of the feature bits.
373        """
374        feature_dps: list[RoborockZeoProtocol] = []
375        if self._feature_bits:
376            feature_dps.extend(build_feature_dp_list(self._feature_bits))
377        if supports_uv_light(self._model):
378            feature_dps.append(RoborockZeoProtocol.UV_LIGHT)
379        if not feature_dps:
380            return
381        try:
382            await self.query_values(feature_dps)
383        except RoborockException as exc:
384            _LOGGER.warning("Feature DPS load failed (non-fatal): %s", exc)
385
386    def supports(self, feature: ZeoFeatureBits) -> bool:
387        """Check whether the device supports a given feature bit."""
388        return bool(self._feature_bits & (1 << feature.value))
389
390    async def query_values(self, protocols: list[RoborockZeoProtocol]) -> dict[RoborockZeoProtocol, Any]:
391        """Query the device for the values of the given protocols."""
392        response = await send_decoded_command(
393            self._channel,
394            {RoborockZeoProtocol.ID_QUERY: protocols},
395            value_encoder=json.dumps,
396        )
397        values = {protocol: convert_zeo_value(protocol, response.get(protocol)) for protocol in protocols}
398        self._merge_query_response(values)
399        return values
400
401    async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[RoborockZeoProtocol, Any]:
402        """Set a value for a specific protocol on the device."""
403        params = {protocol: value}
404        return await send_decoded_command(self._channel, params, value_encoder=lambda x: x)
405
406
407def _parse_device_status(device_status: dict | None) -> dict[int, Any] | None:
408    """Normalize the cloud home data status snapshot to integer datapoint codes."""
409    if not device_status:
410        return None
411    try:
412        return {int(code): value for code, value in device_status.items()}
413    except (TypeError, ValueError):
414        _LOGGER.debug("Ignoring malformed device status snapshot: %s", device_status)
415        return None
416
417
418def create(product: HomeDataProduct, mqtt_channel: MqttChannel, device_status: dict | None = None) -> DyadApi | ZeoApi:
419    """Create traits for A01 devices.
420
421    The optional `device_status` is the cloud home data status snapshot, used
422    to seed `values` so state is available before the first device round trip.
423    """
424    initial_status = _parse_device_status(device_status)
425    match product.category:
426        case RoborockCategory.WET_DRY_VAC:
427            return DyadApi(mqtt_channel, initial_status=initial_status)
428        case RoborockCategory.WASHING_MACHINE:
429            return ZeoApi(mqtt_channel, model=product.model, initial_status=initial_status)
430        case _:
431            raise NotImplementedError(f"Unsupported category {product.category}")