STYJ02YM: Manual movement and mop mode support

miio/viomivacuum.py

@@ -1,10 +1,11 @@
+import click
 import logging
+import time
+
 from collections import defaultdict
 from datetime import timedelta
 from enum import Enum
 
-import click
-
 from .click_common import EnumType, command, format_output
 from .device import Device
 from .utils import pretty_seconds
@@ -29,6 +30,12 @@ class ViomiVacuumState(Enum):
     Docked = 5
 
 
+class ViomiMopMode(Enum):
+    Off = 0  # No Mop, Vacuum only
+    Mixed = 1
+    MopOnly = 2
+
+
 class ViomiLanguage(Enum):
     CN = 1  # Chinese (default)
     EN = 2  # English
@@ -45,6 +52,15 @@ class ViomiCarpetTurbo(Enum):
     Turbo = 2
 
 
+class ViomiMovementDirection(Enum):
+    Forward = 1
+    Left = 2  # Rotate
+    Right = 3  # Rotate
+    Backward = 4
+    Stop = 5
+    # 10 is unknown
+
+
 class ViomiVacuumStatus:
     def __init__(self, data):
         # ["run_state","mode","err_state","battary_life","box_type","mop_type","s_time","s_area",
@@ -140,9 +156,8 @@ def has_new_map(self) -> bool:
         return bool(self.data["has_newmap"])
 
     @property
-    def is_mop(self) -> bool:
-        """True if mopping."""
-        return bool(self.data["is_mop"])
+    def is_mop(self) -> ViomiMopMode:
+        return ViomiMopMode(self.data["is_mop"])
 
 
 class ViomiVacuum(Device):
@@ -206,16 +221,33 @@ def pause(self):
         """Pause cleaning."""
         self.send("set_mode_withroom", [0, 2, 0])
 
-    @command(click.argument("speed", type=str))
-    def set_fan_speed(self, speed: str):
+    @command(click.argument("speed", type=EnumType(ViomiVacuumSpeed, False)))
+    def set_fan_speed(self, speed: ViomiVacuumSpeed):
         """Set fanspeed [silent, standard, medium, turbo]."""
-        self.send("set_suction", [ViomiVacuumSpeed(speed.capitalize()).value])
+        self.send("set_suction", [speed.value])
 
     @command()
     def home(self):
         """Return to home."""
         self.send("set_charge", [1])
 
+    @command(
+        click.argument("direction", type=EnumType(ViomiMovementDirection, False)),
+        click.option('--duration', type=float, default=.5, help='number of seconds to perform this movement'),
+    )
+    def move(self, direction, duration=.5):
+        """Manual movement."""
+        start = time.time()
+        while time.time() - start < duration:
+          self.send("set_direction", [direction.value])
+          time.sleep(.1)
+        self.send("set_direction", [ViomiMovementDirection.Stop.value])
+
+    @command(click.argument("mode", type=EnumType(ViomiMopMode, False)))
+    def mop_mode(self, mode):
+        """Set mopping mode."""
+        self.send("set_mop", [mode.value])
+
     @command()
     def dnd_status(self):
         """Returns do-not-disturb status."""