My shot at 'Vacuum: Provide human readable fan speeds (rytilahti#468)'

miio/vacuum.py

@@ -23,6 +23,22 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+MODELS_VACUUM_V1 = ['rockrobo.vacuum.v1']
+MODELS_VACUUM_V2 = ['roborock.vacuum.s5']
+
+
+FanSpeed = enum.Enum('FanSpeed', [
+        ('Quiet', 'quiet'),
+        ('Balanced', 'balanced'),
+        ('Turbo', 'turbo'),
+        ('Max', 'max'),
+        ('Mop', 'mop'),
+    ]
+    # comment if only levels above are wanted...
+    + [("Level_%d" % x, x) for x in range(0, 101)]
+)
+
+
 class VacuumException(DeviceException):
     pass
 
@@ -32,6 +48,23 @@ class TimerState(enum.Enum):
     Off = "off"
 
 
+@enum.unique
+class _FanSpeedV1(enum.Enum):
+    Quiet = 38
+    Balanced = 60
+    Turbo = 77
+    Max = 90
+
+
+@enum.unique
+class _FanSpeedV2(enum.Enum):
+    Quiet = 101     # -> 38
+    Balanced = 102  # -> 60
+    Turbo = 103     # -> 75
+    Max = 104       # -> 100
+    Mop = 105
+
+
 class Consumable(enum.Enum):
     MainBrush = "main_brush_work_time"
     SideBrush = "side_brush_work_time"
@@ -46,6 +79,32 @@ def __init__(self, ip: str, token: str = None, start_id: int = 0,
                  debug: int = 0) -> None:
         super().__init__(ip, token, start_id, debug)
         self.manual_seqnum = -1
+        self._model = self.info().model
+        self._init_fan_speeds()
+
+    def _init_fan_speeds(self):
+        # initialize int -> enum fan speed map
+
+        if not hasattr(self, '_fanspeedMap'):
+            fanspeedMap = dict()
+            for fs in list(FanSpeed):
+                if type(fs.value) is int:
+                    fanspeedMap[fs.value] = FanSpeed(fs.value)
+
+            if self._model in MODELS_VACUUM_V1:
+                for fs in list(_FanSpeedV1):
+                    fanspeedMap[fs.value] = FanSpeed(fs.name.lower())
+
+            elif self._model in MODELS_VACUUM_V2:
+                _cur = self.fan_speed()
+
+                for fs in list(_FanSpeedV2):
+                    self.set_fan_speed(fs.value)
+                    fanspeedMap[self.fan_speed()] = FanSpeed(fs.name.lower())
+
+                self.set_fan_speed(_cur)
+
+            self._fanspeedMap = fanspeedMap
 
     @command()
     def start(self):
@@ -412,19 +471,50 @@ def disable_dnd(self):
         return self.send("close_dnd_timer", [""])
 
     @command(
-        click.argument("speed", type=int),
+        click.argument("speed", type=str),
     )
-    def set_fan_speed(self, speed: int):
+    def set_fan_speed(self, speed: [int, FanSpeed]):
         """Set fan speed.
 
         :param int speed: Fan speed to set"""
         # speed = [38, 60 or 77]
+        if isinstance(speed, FanSpeed):
+            try:
+                speed = int(speed.value)
+
+            except ValueError:
+                for k, v in self._fanspeedMap.items():
+                    if v == speed:
+                        speed = k
+                        break
+
         return self.send("set_custom_mode", [speed])
 
     @command()
     def fan_speed(self):
         """Return fan speed."""
-        return self.send("get_custom_mode")[0]
+        speed = self.send("get_custom_mode")[0]
+        try:
+            if hasattr(self, '_fanspeedMap'):
+                return self._fanspeedMap[speed]
+        except KeyError:
+            pass
+
+        return speed
+
+    @command()
+    def fan_speed_list(self):
+        """Get the list of available fan speed steps of the vacuum cleaner."""
+
+        if not hasattr(self, '_fanspeeds'):
+            fanspeeds = list(FanSpeed)
+
+            if self._model in MODELS_VACUUM_V1:
+                fanspeeds.remove(FanSpeed.Mop)
+
+            self._fanspeeds = fanspeeds
+
+        return self._fanspeeds
 
     @command()
     def sound_info(self):

miio/vacuum_cli.py

@@ -314,17 +314,31 @@ def dnd(vac: miio.Vacuum, cmd: str,
 
 
 @cli.command()
-@click.argument('speed', type=int, required=False)
+@click.argument('speed', type=str, required=False)
 @pass_dev
 def fanspeed(vac: miio.Vacuum, speed):
     """Query and adjust the fan speed."""
+    from miio.vacuum import FanSpeed
     if speed:
         click.echo("Setting fan speed to %s" % speed)
+        try:
+            speed = int(speed)
+        except ValueError:
+            speed = FanSpeed(speed.lower())
+
         vac.set_fan_speed(speed)
     else:
         click.echo("Current fan speed: %s" % vac.fan_speed())
 
 
+@cli.command()
+@pass_dev
+def fanspeed_list(vac: miio.Vacuum):
+    """Query the list of available fan speed steps of the vacuum cleaner."""
+
+    click.echo("%s" % [e.value for e in vac.fan_speed_list()])
+
+
 @cli.group(invoke_without_command=True)
 @pass_dev
 @click.pass_context