Improve base station action

DeebotUniverse · Dec 20, 2024 · b3e456e · b3e456e
1 parent 260db3f
commit b3e456e
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 10 deletions.
diff --git a/deebot_client/capabilities.py b/deebot_client/capabilities.py
@@ -109,10 +109,10 @@ class CapabilitySetEnable(CapabilitySet[_EVENT, [bool]]):
 
 
 @dataclass(frozen=True)
-class CapabilityExecute:
+class CapabilityExecute(Generic[_P]):
     """Capability to execute a command."""
 
-    execute: type[Command]
+    execute: Callable[_P, Command]
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -122,6 +122,11 @@ class CapabilityTypes(Generic[_T]):
     types: tuple[_T, ...]
 
 
+@dataclass(frozen=True, kw_only=True)
+class CapabilityExecuteTypes(CapabilityTypes[_T], CapabilityExecute[[_T]]):
+    """Capability to execute a command with types."""
+
+
 @dataclass(frozen=True, kw_only=True)
 class CapabilitySetTypes(CapabilitySet[_EVENT, _P], CapabilityTypes[_T]):
     """Capability for set command and types."""
@@ -169,11 +174,11 @@ class CapabilityMap:
 
     cached_info: CapabilityEvent[CachedMapInfoEvent]
     changed: CapabilityEvent[MapChangedEvent]
-    clear: CapabilityExecute | None = None
+    clear: CapabilityExecute[[]] | None = None
     major: CapabilityEvent[MajorMapEvent]
     multi_state: CapabilitySetEnable[MultimapStateEvent]
     position: CapabilityEvent[PositionsEvent]
-    relocation: CapabilityExecute
+    relocation: CapabilityExecute[[]]
     rooms: CapabilityEvent[RoomsEvent]
     trace: CapabilityEvent[MapTraceEvent]
 
@@ -216,7 +221,7 @@ class CapabilitySettings:
 class CapabilityBaseStation:
     """Capabilities for the base station."""
 
-    action: Callable[[BaseStationAction], Command]
+    action: CapabilityExecuteTypes[BaseStationAction]
     auto_empty: CapabilitySetTypes[
         auto_empty.AutoEmptyEvent,
         [bool | None, auto_empty.Frequency | str | None],
@@ -234,7 +239,7 @@ class Capabilities(ABC):
     availability: CapabilityEvent[AvailabilityEvent]
     base_station: CapabilityBaseStation | None = None
     battery: CapabilityEvent[BatteryEvent]
-    charge: CapabilityExecute
+    charge: CapabilityExecute[[]]
     clean: CapabilityClean
     custom: CapabilityCustomCommand[CustomCommandEvent]
     error: CapabilityEvent[ErrorEvent]
@@ -244,7 +249,7 @@ class Capabilities(ABC):
     life_span: CapabilityLifeSpan
     map: CapabilityMap | None = None
     network: CapabilityEvent[NetworkInfoEvent]
-    play_sound: CapabilityExecute
+    play_sound: CapabilityExecute[[]]
     settings: CapabilitySettings
     state: CapabilityEvent[StateEvent]
     stats: CapabilityStats

diff --git a/deebot_client/hardware/deebot/buom7k.py b/deebot_client/hardware/deebot/buom7k.py
@@ -10,6 +10,7 @@
     CapabilityCustomCommand,
     CapabilityEvent,
     CapabilityExecute,
+    CapabilityExecuteTypes,
     CapabilityLifeSpan,
     CapabilityMap,
     CapabilitySet,
@@ -19,6 +20,7 @@
     CapabilityStats,
     DeviceType,
 )
+from deebot_client.commands import BaseStationAction
 from deebot_client.commands.json.auto_empty import GetAutoEmpty, SetAutoEmpty
 from deebot_client.commands.json.battery import GetBattery
 from deebot_client.commands.json.carpet import (
@@ -100,7 +102,9 @@
             AvailabilityEvent, [GetBattery(is_available_check=True)]
         ),
         base_station=CapabilityBaseStation(
-            action=StationAction,
+            action=CapabilityExecuteTypes(
+                StationAction, types=(BaseStationAction.EMPTY_DUSTBIN,)
+            ),
             auto_empty=CapabilitySetTypes(
                 event=AutoEmptyEvent,
                 get=[GetAutoEmpty()],

diff --git a/deebot_client/hardware/deebot/p95mgv.py b/deebot_client/hardware/deebot/p95mgv.py
@@ -10,6 +10,7 @@
     CapabilityCustomCommand,
     CapabilityEvent,
     CapabilityExecute,
+    CapabilityExecuteTypes,
     CapabilityLifeSpan,
     CapabilityMap,
     CapabilitySet,
@@ -19,6 +20,7 @@
     CapabilityStats,
     DeviceType,
 )
+from deebot_client.commands import BaseStationAction
 from deebot_client.commands.json.advanced_mode import GetAdvancedMode, SetAdvancedMode
 from deebot_client.commands.json.auto_empty import GetAutoEmpty, SetAutoEmpty
 from deebot_client.commands.json.battery import GetBattery
@@ -118,7 +120,9 @@
             AvailabilityEvent, [GetBattery(is_available_check=True)]
         ),
         base_station=CapabilityBaseStation(
-            action=StationAction,
+            action=CapabilityExecuteTypes(
+                StationAction, types=(BaseStationAction.EMPTY_DUSTBIN,)
+            ),
             auto_empty=CapabilitySetTypes(
                 event=auto_empty.AutoEmptyEvent,
                 get=[GetAutoEmpty()],

diff --git a/deebot_client/hardware/deebot/qhe2o2.py b/deebot_client/hardware/deebot/qhe2o2.py
@@ -10,6 +10,7 @@
     CapabilityCustomCommand,
     CapabilityEvent,
     CapabilityExecute,
+    CapabilityExecuteTypes,
     CapabilityLifeSpan,
     CapabilityMap,
     CapabilitySet,
@@ -19,6 +20,7 @@
     CapabilityStats,
     DeviceType,
 )
+from deebot_client.commands import BaseStationAction
 from deebot_client.commands.json.auto_empty import GetAutoEmpty, SetAutoEmpty
 from deebot_client.commands.json.battery import GetBattery
 from deebot_client.commands.json.carpet import (
@@ -101,7 +103,9 @@
         ),
         battery=CapabilityEvent(BatteryEvent, [GetBattery()]),
         base_station=CapabilityBaseStation(
-            action=StationAction,
+            action=CapabilityExecuteTypes(
+                StationAction, types=(BaseStationAction.EMPTY_DUSTBIN,)
+            ),
             auto_empty=CapabilitySetTypes(
                 event=AutoEmptyEvent,
                 get=[GetAutoEmpty()],