[
  {
    "path": ".gitignore",
    "content": "__pycache__\n"
  },
  {
    "path": "README.md",
    "content": "# Hacky Home assistant support for Xiaomi vacuum STYJ02YM \n\n### Install:\n- Create the following folder structure: /config/custom_components/miio2 and place all files there [4 files](https://github.com/nqkdev/home-assistant-vacuum-styj02ym) there.\n- Add the configuration to configuration.yaml, example:\n\n```yaml\nvacuum:\n  - platform: miio2\n    host: 192.168.68.105\n    token: !secret vacuum\n    name: Mi hihi\n```\n"
  },
  {
    "path": "__init__.py",
    "content": "\n"
  },
  {
    "path": "manifest.json",
    "content": "{\n    \"domain\": \"miio2\",\n    \"name\": \"Xiaomi miio vacuum STYJ02YM\",\n    \"version\": \"1.0.0\",\n    \"documentation\": \"\",\n    \"requirements\": [\n        \"construct==2.10.56\",\n        \"python-miio>=0.5.12\"\n    ],\n    \"dependencies\": [],\n    \"codeowners\": []\n}\n"
  },
  {
    "path": "vacuum.py",
    "content": "\"\"\"Support for the Xiaomi vacuum cleaner robot.\"\"\"\nimport asyncio\nfrom functools import partial\nimport logging\n\nfrom miio import DeviceException, RoborockVacuum  # pylint: disable=import-error\nimport voluptuous as vol\n\nfrom homeassistant.components.vacuum import (\n    ATTR_CLEANED_AREA,\n    DOMAIN,\n    PLATFORM_SCHEMA,\n    STATE_CLEANING,\n    STATE_DOCKED,\n    STATE_ERROR,\n    STATE_IDLE,\n    STATE_PAUSED,\n    STATE_RETURNING,\n    SUPPORT_BATTERY,\n    SUPPORT_FAN_SPEED,\n    SUPPORT_LOCATE,\n    SUPPORT_PAUSE,\n    SUPPORT_RETURN_HOME,\n    SUPPORT_SEND_COMMAND,\n    SUPPORT_START,\n    SUPPORT_STATE,\n    SUPPORT_STOP,\n    StateVacuumEntity,\n)\nfrom homeassistant.const import (\n    ATTR_ENTITY_ID,\n    CONF_HOST,\n    CONF_NAME,\n    CONF_TOKEN,\n    STATE_OFF,\n    STATE_ON,\n)\nimport homeassistant.helpers.config_validation as cv\n\n_LOGGER = logging.getLogger(__name__)\n\nDEFAULT_NAME = \"Xiaomi Vacuum cleaner STYJ02YM\"\nDATA_KEY = \"vacuum.miio2\"\n\nPLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(\n    {\n        vol.Required(CONF_HOST): cv.string,\n        vol.Required(CONF_TOKEN): vol.All(str, vol.Length(min=32, max=32)),\n        vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,\n    },\n    extra=vol.ALLOW_EXTRA,\n)\n\nVACUUM_SERVICE_SCHEMA = vol.Schema({vol.Optional(ATTR_ENTITY_ID): cv.comp_entity_ids})\nSERVICE_CLEAN_ZONE = \"xiaomi_clean_zone\"\nSERVICE_CLEAN_POINT = \"xiaomi_clean_point\"\nATTR_ZONE_ARRAY = \"zone\"\nATTR_ZONE_REPEATER = \"repeats\"\nATTR_POINT = \"point\"\nSERVICE_SCHEMA_CLEAN_ZONE = VACUUM_SERVICE_SCHEMA.extend(\n    {\n        vol.Required(ATTR_ZONE_ARRAY): vol.All(\n            list,\n            [\n                vol.ExactSequence(\n                    [vol.Coerce(float), vol.Coerce(float), vol.Coerce(float), vol.Coerce(float)]\n                )\n            ],\n        ),\n        vol.Required(ATTR_ZONE_REPEATER): vol.All(\n            vol.Coerce(int), vol.Clamp(min=1, max=3)\n        ),\n    }\n)\nSERVICE_SCHEMA_CLEAN_POINT = VACUUM_SERVICE_SCHEMA.extend(\n    {\n        vol.Required(ATTR_POINT): vol.All(\n            vol.ExactSequence(\n                [vol.Coerce(float), vol.Coerce(float)]\n            )\n        )\n    }\n)\nSERVICE_TO_METHOD = {\n    SERVICE_CLEAN_ZONE: {\n        \"method\": \"async_clean_zone\",\n        \"schema\": SERVICE_SCHEMA_CLEAN_ZONE,\n    },\n    SERVICE_CLEAN_POINT: {\n        \"method\": \"async_clean_point\",\n        \"schema\": SERVICE_SCHEMA_CLEAN_POINT,\n    }\n}\n\nFAN_SPEEDS = {\"Silent\": 0, \"Standard\": 1, \"Medium\": 2, \"Turbo\": 3}\n\n\nSUPPORT_XIAOMI = (\n    SUPPORT_STATE\n    | SUPPORT_PAUSE\n    | SUPPORT_STOP\n    | SUPPORT_RETURN_HOME\n    | SUPPORT_FAN_SPEED\n    | SUPPORT_LOCATE\n    | SUPPORT_SEND_COMMAND\n    | SUPPORT_BATTERY\n    | SUPPORT_START\n)\n\n\nSTATE_CODE_TO_STATE = {\n    0: STATE_IDLE,\n    1: STATE_IDLE,\n    2: STATE_PAUSED,\n    3: STATE_CLEANING,\n    4: STATE_RETURNING,\n    5: STATE_DOCKED,\n    6: STATE_CLEANING,  # Vacuum & Mop\n    7: STATE_CLEANING   # Mop only\n}\n\nALL_PROPS = [\"run_state\", \"mode\", \"err_state\", \"battary_life\", \"box_type\", \"mop_type\", \"s_time\",\n             \"s_area\", \"suction_grade\", \"water_grade\", \"remember_map\", \"has_map\", \"is_mop\", \"has_newmap\"]\n\n\nasync def async_setup_platform(hass, config, async_add_entities, discovery_info=None):\n  \"\"\"Set up the Xiaomi vacuum cleaner robot platform.\"\"\"\n  if DATA_KEY not in hass.data:\n    hass.data[DATA_KEY] = {}\n\n  host = config[CONF_HOST]\n  token = config[CONF_TOKEN]\n  name = config[CONF_NAME]\n\n  # Create handler\n  _LOGGER.info(\"Initializing with host %s (token %s...)\", host, token[:5])\n  vacuum = RoborockVacuum(host, token)\n\n  mirobo = MiroboVacuum2(name, vacuum)\n  hass.data[DATA_KEY][host] = mirobo\n\n  async_add_entities([mirobo], update_before_add=True)\n\n  async def async_service_handler(service):\n    \"\"\"Map services to methods on MiroboVacuum.\"\"\"\n    method = SERVICE_TO_METHOD.get(service.service)\n    params = {\n        key: value for key, value in service.data.items() if key != ATTR_ENTITY_ID\n    }\n    entity_ids = service.data.get(ATTR_ENTITY_ID)\n\n    if entity_ids:\n      target_vacuums = [\n          vac\n          for vac in hass.data[DATA_KEY].values()\n          if vac.entity_id in entity_ids\n      ]\n    else:\n      target_vacuums = hass.data[DATA_KEY].values()\n\n    update_tasks = []\n    for vacuum in target_vacuums:\n      await getattr(vacuum, method[\"method\"])(**params)\n\n    for vacuum in target_vacuums:\n      update_coro = vacuum.async_update_ha_state(True)\n      update_tasks.append(update_coro)\n\n    if update_tasks:\n      await asyncio.wait(update_tasks)\n\n  for vacuum_service in SERVICE_TO_METHOD:\n    schema = SERVICE_TO_METHOD[vacuum_service].get(\"schema\", VACUUM_SERVICE_SCHEMA)\n    hass.services.async_register(\n        DOMAIN, vacuum_service, async_service_handler, schema=schema\n    )\n\n\nclass MiroboVacuum2(StateVacuumEntity):\n  \"\"\"Representation of a Xiaomi Vacuum cleaner robot.\"\"\"\n\n  def __init__(self, name, vacuum):\n    \"\"\"Initialize the Xiaomi vacuum cleaner robot handler.\"\"\"\n    self._name = name\n    self._vacuum = vacuum\n\n    self._last_clean_point = None\n\n    self.vacuum_state = None\n    self._available = False\n\n  @property\n  def name(self):\n    \"\"\"Return the name of the device.\"\"\"\n    return self._name\n\n  @property\n  def state(self):\n    \"\"\"Return the status of the vacuum cleaner.\"\"\"\n    if self.vacuum_state is not None:\n      # The vacuum reverts back to an idle state after erroring out.\n      # We want to keep returning an error until it has been cleared.\n\n      try:\n        return STATE_CODE_TO_STATE[int(self.vacuum_state['run_state'])]\n      except KeyError:\n        _LOGGER.error(\n            \"STATE not supported, state_code: %s\",\n            self.vacuum_state['run_state'],\n        )\n        return None\n\n  @property\n  def battery_level(self):\n    \"\"\"Return the battery level of the vacuum cleaner.\"\"\"\n    if self.vacuum_state is not None:\n      return self.vacuum_state['battary_life']\n\n  @property\n  def fan_speed(self):\n    \"\"\"Return the fan speed of the vacuum cleaner.\"\"\"\n    if self.vacuum_state is not None:\n      speed = self.vacuum_state['suction_grade']\n      if speed in FAN_SPEEDS.values():\n        return [key for key, value in FAN_SPEEDS.items() if value == speed][0]\n      return speed\n\n  @property\n  def fan_speed_list(self):\n    \"\"\"Get the list of available fan speed steps of the vacuum cleaner.\"\"\"\n    return list(sorted(FAN_SPEEDS.keys(), key=lambda s: FAN_SPEEDS[s]))\n\n  @property\n  def device_state_attributes(self):\n    \"\"\"Return the specific state attributes of this vacuum cleaner.\"\"\"\n    attrs = {}\n    if self.vacuum_state is not None:\n      attrs.update(self.vacuum_state)\n      try:\n        attrs['status'] = STATE_CODE_TO_STATE[int(self.vacuum_state['run_state'])]\n      except KeyError:\n        return \"Definition missing for state %s\" % self.vacuum_state['run_state']\n    return attrs\n\n  @property\n  def available(self) -> bool:\n    \"\"\"Return True if entity is available.\"\"\"\n    return self._available\n\n  @property\n  def supported_features(self):\n    \"\"\"Flag vacuum cleaner robot features that are supported.\"\"\"\n    return SUPPORT_XIAOMI\n\n  async def _try_command(self, mask_error, func, *args, **kwargs):\n    \"\"\"Call a vacuum command handling error messages.\"\"\"\n    try:\n      await self.hass.async_add_executor_job(partial(func, *args, **kwargs))\n      return True\n    except DeviceException as exc:\n      _LOGGER.error(mask_error, exc)\n      return False\n\n  async def async_start(self):\n    \"\"\"Start or resume the cleaning task.\"\"\"\n    mode = self.vacuum_state['mode']\n    is_mop = self.vacuum_state['is_mop']\n    actionMode = 0\n\n    if mode == 4 and self._last_clean_point is not None:\n      method = 'set_pointclean'\n      param = [1, self._last_clean_point[0], self._last_clean_point[1]]\n    else:\n      if mode == 2:\n        actionMode = 2\n      else:\n        if is_mop == 2:\n          actionMode = 3\n        else:\n          actionMode = is_mop\n      if mode == 3:\n        method = 'set_mode'\n        param = [3, 1]\n      else:\n        method = 'set_mode_withroom'\n        param = [actionMode, 1, 0]\n    await self._try_command(\"Unable to start the vacuum: %s\", self._vacuum.raw_command, method, param)\n\n  async def async_pause(self):\n    \"\"\"Pause the cleaning task.\"\"\"\n    mode = self.vacuum_state['mode']\n    is_mop = self.vacuum_state['is_mop']\n    actionMode = 0\n\n    if mode == 4 and self._last_clean_point is not None:\n      method = 'set_pointclean'\n      param = [3, self._last_clean_point[0], self._last_clean_point[1]]\n    else:\n      if mode == 2:\n        actionMode = 2\n      else:\n        if is_mop == 2:\n          actionMode = 3\n        else:\n          actionMode = is_mop\n      if mode == 3:\n        method = 'set_mode'\n        param = [3, 3]\n      else:\n        method = 'set_mode_withroom'\n        param = [actionMode, 3, 0]\n    await self._try_command(\"Unable to set pause: %s\", self._vacuum.raw_command, method, param)\n\n  async def async_stop(self, **kwargs):\n    \"\"\"Stop the vacuum cleaner.\"\"\"\n    mode = self.vacuum_state['mode']\n    if mode == 3:\n      method = 'set_mode'\n      param = [3, 0]\n    elif mode == 4:\n      method = 'set_pointclean'\n      param = [0, 0, 0]\n      self._last_clean_point = None\n    else:\n      method = 'set_mode'\n      param = [0]\n    await self._try_command(\"Unable to stop: %s\", self._vacuum.raw_command, method, param)\n\n  async def async_set_fan_speed(self, fan_speed, **kwargs):\n    \"\"\"Set fan speed.\"\"\"\n    if fan_speed.capitalize() in FAN_SPEEDS:\n      fan_speed = FAN_SPEEDS[fan_speed.capitalize()]\n    else:\n      try:\n        fan_speed = int(fan_speed)\n      except ValueError as exc:\n        _LOGGER.error(\n            \"Fan speed step not recognized (%s). \" \"Valid speeds are: %s\",\n            exc,\n            self.fan_speed_list,\n        )\n        return\n    await self._try_command(\n        \"Unable to set fan speed: %s\", self._vacuum.raw_command, 'set_suction', [\n            fan_speed]\n    )\n\n  async def async_return_to_base(self, **kwargs):\n    \"\"\"Set the vacuum cleaner to return to the dock.\"\"\"\n    await self._try_command(\"Unable to return home: %s\", self._vacuum.raw_command, 'set_charge', [1])\n\n  async def async_locate(self, **kwargs):\n    \"\"\"Locate the vacuum cleaner.\"\"\"\n    await self._try_command(\"Unable to locate the botvac: %s\", self._vacuum.raw_command, 'set_resetpos', [1])\n\n  async def async_send_command(self, command, params=None, **kwargs):\n    \"\"\"Send raw command.\"\"\"\n    await self._try_command(\n        \"Unable to send command to the vacuum: %s\",\n        self._vacuum.raw_command,\n        command,\n        params,\n    )\n\n  def update(self):\n    \"\"\"Fetch state from the device.\"\"\"\n    try:\n      state = self._vacuum.raw_command('get_prop', ALL_PROPS)\n\n      self.vacuum_state = dict(zip(ALL_PROPS, state))\n\n      self._available = True\n      \n      # Automatically set mop based on mop_type\n      is_mop = bool(self.vacuum_state['is_mop'])\n      has_mop = bool(self.vacuum_state['mop_type'])\n\n      update_mop = None\n      if is_mop and not has_mop:\n        update_mop = 0\n      elif not is_mop and has_mop:\n        update_mop = 1\n\n      if update_mop is not None:\n        self._vacuum.raw_command('set_mop', [update_mop])\n        self.update()\n    except OSError as exc:\n      _LOGGER.error(\"Got OSError while fetching the state: %s\", exc)\n    except DeviceException as exc:\n      _LOGGER.warning(\"Got exception while fetching the state: %s\", exc)\n\n  async def async_clean_zone(self, zone, repeats=1):\n    \"\"\"Clean selected area for the number of repeats indicated.\"\"\"\n    result = []\n    i = 0\n    for z in zone:\n      x1, y2, x2, y1 = z\n      res = '_'.join(str(x) for x in [i, 0, x1, y1, x1, y2, x2, y2, x2, y1])\n      for _ in range(repeats):\n        result.append(res)\n        i += 1\n    result = [i] + result\n\n    await self._try_command(\"Unable to clean zone: %s\", self._vacuum.raw_command, 'set_uploadmap', [1]) \\\n        and await self._try_command(\"Unable to clean zone: %s\", self._vacuum.raw_command, 'set_zone', result) \\\n        and await self._try_command(\"Unable to clean zone: %s\", self._vacuum.raw_command, 'set_mode', [3, 1])\n\n  async def async_clean_point(self, point):\n    \"\"\"Clean selected area\"\"\"\n    x, y = point\n    self._last_clean_point = point\n    await self._try_command(\"Unable to clean point: %s\", self._vacuum.raw_command, 'set_uploadmap', [0]) \\\n        and await self._try_command(\"Unable to clean point: %s\", self._vacuum.raw_command, 'set_pointclean', [1, x, y])\n"
  }
]