Add water capability

Author: nanomad

deebot_client/commands/xml/water_info.py

@@ -2,15 +2,18 @@
 
 from __future__ import annotations
 
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any
 
+from deebot_client.command import InitParam
 from deebot_client.events import (
     WaterAmount,
     WaterInfoEvent,
 )
 from deebot_client.message import HandlingResult
+from deebot_client.util import get_enum
 
-from .common import XmlGetCommand
+from .common import XmlGetCommand, XmlSetCommand
 
 if TYPE_CHECKING:
     from xml.etree.ElementTree import Element
@@ -47,6 +50,23 @@ def _handle_xml(cls, event_bus: EventBus, xml: Element) -> HandlingResult:
         return HandlingResult.success()
 
 
+class SetWaterPermeability(XmlSetCommand):
+    """SetWaterPermeability command."""
+
+    NAME = "SetWaterPermeability"
+    get_command = GetWaterPermeability
+    _mqtt_params = MappingProxyType(
+        {
+            "amount": InitParam(WaterAmount),
+        }
+    )
+
+    def __init__(self, amount: WaterAmount | str) -> None:
+        if isinstance(amount, str):
+            amount = get_enum(WaterAmount, amount)
+        super().__init__({"v": str(amount.value)})
+
+
 class GetWaterBoxInfo(XmlGetCommand):
     """GetWaterBoxInfo command."""
 
@@ -60,11 +80,7 @@ def handle_set_args(
 
         :return: A message response
         """
-        event_bus.notify(
-            WaterInfoEvent(
-                amount=WaterAmount.HIGH, mop_attached=(str(args["on"]) != "0")
-            )
-        )
+        event_bus.notify(WaterInfoEvent(mop_attached=(str(args["on"]) != "0")))
         return HandlingResult.success()
 
     @classmethod
@@ -76,7 +92,5 @@ def _handle_xml(cls, event_bus: EventBus, xml: Element) -> HandlingResult:
         if xml.attrib.get("ret") != "ok" or not (on := xml.attrib.get("on")):
             return HandlingResult.analyse()
 
-        event_bus.notify(
-            WaterInfoEvent(amount=WaterAmount.HIGH, mop_attached=on != "0")
-        )
+        event_bus.notify(WaterInfoEvent(mop_attached=on != "0"))
         return HandlingResult.success()

deebot_client/events/water_info.py

@@ -30,7 +30,7 @@ class SweepType(IntEnum):
 class WaterInfoEvent(Event):
     """Water info event representation."""
 
-    amount: WaterAmount
     # None means no data available
+    amount: WaterAmount | None = None
     sweep_type: SweepType | None = None
     mop_attached: bool | None = field(kw_only=True, default=None)

deebot_client/hardware/deebot/2pv572.py

@@ -26,12 +26,14 @@
     GetCleanSpeed,
     GetCleanState,
     GetLifeSpan,
+    GetWaterPermeability,
     PlaySound,
     SetCleanSpeed,
 )
 from deebot_client.commands.xml.charge_state import GetChargeState
 from deebot_client.commands.xml.error import GetError
 from deebot_client.commands.xml.stats import GetCleanSum
+from deebot_client.commands.xml.water_info import GetWaterBoxInfo, SetWaterPermeability
 from deebot_client.const import DataType
 from deebot_client.events import (
     AvailabilityEvent,
@@ -48,6 +50,8 @@
     StateEvent,
     StatsEvent,
     TotalStatsEvent,
+    WaterAmount,
+    WaterInfoEvent,
 )
 from deebot_client.models import StaticDeviceInfo
 from deebot_client.util import short_name
@@ -97,5 +101,15 @@
             total=CapabilityEvent(TotalStatsEvent, [GetCleanSum()]),
         ),
         settings=CapabilitySettings(),
+        water=CapabilitySetTypes(
+            event=WaterInfoEvent,
+            get=[GetWaterPermeability(), GetWaterBoxInfo()],
+            set=SetWaterPermeability,
+            types=(
+                WaterAmount.LOW,
+                WaterAmount.MEDIUM,
+                WaterAmount.HIGH,
+            ),
+        ),
     ),
 )