Add custom and zone cleaning to Neato Vacuums (#20779)

* Adding custom and zone cleaning to Neato Vacuums

* Fixing line length and missing imports

* Line too long

* Adding details to the custom service

* Fix linting issues

* Reverting ACTION

* Code cleanup

* Typo

* Requested modifications

* Changing the custom service domain

* No service schema depency anymore

* Removing useless code

* Linting

* Requested changes

* Requested changes for domain

* Revert the service domain back to vacuum

homeassistant/components/neato/__init__.py

@@ -18,6 +18,7 @@ DOMAIN = 'neato'
 NEATO_ROBOTS = 'neato_robots'
 NEATO_LOGIN = 'neato_login'
 NEATO_MAP_DATA = 'neato_map_data'
+NEATO_PERSISTENT_MAPS = 'neato_persistent_maps'
 
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
@@ -197,6 +198,7 @@ class NeatoHub:
             domain_config[CONF_USERNAME],
             domain_config[CONF_PASSWORD])
         self._hass.data[NEATO_ROBOTS] = self.my_neato.robots
+        self._hass.data[NEATO_PERSISTENT_MAPS] = self.my_neato.persistent_maps
         self._hass.data[NEATO_MAP_DATA] = self.my_neato.maps
 
     def login(self):
@@ -216,6 +218,7 @@ class NeatoHub:
         _LOGGER.debug("Running HUB.update_robots %s",
                       self._hass.data[NEATO_ROBOTS])
         self._hass.data[NEATO_ROBOTS] = self.my_neato.robots
+        self._hass.data[NEATO_PERSISTENT_MAPS] = self.my_neato.persistent_maps
         self._hass.data[NEATO_MAP_DATA] = self.my_neato.maps
 
     def download_map(self, url):

homeassistant/components/neato/vacuum.py

@@ -2,15 +2,21 @@
 import logging
 from datetime import timedelta
 import requests
+import voluptuous as vol
 
+from homeassistant.const import (ATTR_ENTITY_ID)
 from homeassistant.components.vacuum import (
     StateVacuumDevice, SUPPORT_BATTERY, SUPPORT_PAUSE, SUPPORT_RETURN_HOME,
     SUPPORT_STATE, SUPPORT_STOP, SUPPORT_START, STATE_IDLE,
     STATE_PAUSED, STATE_CLEANING, STATE_DOCKED, STATE_RETURNING, STATE_ERROR,
     SUPPORT_MAP, ATTR_STATUS, ATTR_BATTERY_LEVEL, ATTR_BATTERY_ICON,
-    SUPPORT_LOCATE, SUPPORT_CLEAN_SPOT)
+    SUPPORT_LOCATE, SUPPORT_CLEAN_SPOT, DOMAIN)
 from homeassistant.components.neato import (
-    NEATO_ROBOTS, NEATO_LOGIN, NEATO_MAP_DATA, ACTION, ERRORS, MODE, ALERTS)
+    NEATO_ROBOTS, NEATO_LOGIN, NEATO_MAP_DATA, ACTION, ERRORS, MODE, ALERTS,
+    NEATO_PERSISTENT_MAPS)
+
+from homeassistant.helpers.service import extract_entity_ids
+import homeassistant.helpers.config_validation as cv
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -19,8 +25,8 @@ DEPENDENCIES = ['neato']
 SCAN_INTERVAL = timedelta(minutes=5)
 
 SUPPORT_NEATO = SUPPORT_BATTERY | SUPPORT_PAUSE | SUPPORT_RETURN_HOME | \
-                 SUPPORT_STOP | SUPPORT_START | SUPPORT_CLEAN_SPOT | \
+    SUPPORT_STOP | SUPPORT_START | SUPPORT_CLEAN_SPOT | \
-                 SUPPORT_STATE | SUPPORT_MAP | SUPPORT_LOCATE
+    SUPPORT_STATE | SUPPORT_MAP | SUPPORT_LOCATE
 
 ATTR_CLEAN_START = 'clean_start'
 ATTR_CLEAN_STOP = 'clean_stop'
@@ -30,15 +36,56 @@ ATTR_CLEAN_BATTERY_END = 'battery_level_at_clean_end'
 ATTR_CLEAN_SUSP_COUNT = 'clean_suspension_count'
 ATTR_CLEAN_SUSP_TIME = 'clean_suspension_time'
 
+ATTR_MODE = 'mode'
+ATTR_NAVIGATION = 'navigation'
+ATTR_CATEGORY = 'category'
+ATTR_ZONE = 'zone'
+
+SERVICE_NEATO_CUSTOM_CLEANING = 'neato_custom_cleaning'
+
+SERVICE_NEATO_CUSTOM_CLEANING_SCHEMA = vol.Schema({
+    vol.Required(ATTR_ENTITY_ID): cv.entity_ids,
+    vol.Optional(ATTR_MODE, default=2): cv.positive_int,
+    vol.Optional(ATTR_NAVIGATION, default=1): cv.positive_int,
+    vol.Optional(ATTR_CATEGORY, default=4): cv.positive_int,
+    vol.Optional(ATTR_ZONE): cv.string
+})
+
 
 def setup_platform(hass, config, add_entities, discovery_info=None):
     """Set up the Neato vacuum."""
     dev = []
     for robot in hass.data[NEATO_ROBOTS]:
         dev.append(NeatoConnectedVacuum(hass, robot))
+
+    if not dev:
+        return
+
     _LOGGER.debug("Adding vacuums %s", dev)
     add_entities(dev, True)
 
+    def neato_custom_cleaning_service(call):
+        """Zone cleaning service that allows user to change options."""
+        for robot in service_to_entities(call):
+            if call.service == SERVICE_NEATO_CUSTOM_CLEANING:
+                mode = call.data.get(ATTR_MODE)
+                navigation = call.data.get(ATTR_NAVIGATION)
+                category = call.data.get(ATTR_CATEGORY)
+                zone = call.data.get(ATTR_ZONE)
+                robot.neato_custom_cleaning(
+                    mode, navigation, category, zone)
+
+    def service_to_entities(call):
+        """Return the known devices that a service call mentions."""
+        entity_ids = extract_entity_ids(hass, call)
+        entities = [entity for entity in dev
+                    if entity.entity_id in entity_ids]
+        return entities
+
+    hass.services.register(DOMAIN, SERVICE_NEATO_CUSTOM_CLEANING,
+                           neato_custom_cleaning_service,
+                           schema=SERVICE_NEATO_CUSTOM_CLEANING_SCHEMA)
+
 
 class NeatoConnectedVacuum(StateVacuumDevice):
     """Representation of a Neato Connected Vacuum."""
@@ -62,6 +109,9 @@ class NeatoConnectedVacuum(StateVacuumDevice):
         self._available = False
         self._battery_level = None
         self._robot_serial = self.robot.serial
+        self._robot_maps = hass.data[NEATO_PERSISTENT_MAPS]
+        self._robot_boundaries = {}
+        self._robot_has_map = self.robot.has_persistent_maps
 
     def update(self):
         """Update the states of Neato Vacuums."""
@@ -129,12 +179,18 @@ class NeatoConnectedVacuum(StateVacuumDevice):
             ['time_in_suspended_cleaning'])
         self.clean_battery_start = (
             self._mapdata[self._robot_serial]['maps'][0]['run_charge_at_start']
-            )
+        )
         self.clean_battery_end = (
             self._mapdata[self._robot_serial]['maps'][0]['run_charge_at_end'])
 
         self._battery_level = self._state['details']['charge']
 
+        if self._robot_has_map:
+            robot_map_id = self._robot_maps[self._robot_serial][0]['id']
+
+            self._robot_boundaries = self.robot.get_map_boundaries(
+                robot_map_id).json()
+
     @property
     def name(self):
         """Return the name of the device."""
@@ -224,3 +280,20 @@ class NeatoConnectedVacuum(StateVacuumDevice):
     def clean_spot(self, **kwargs):
         """Run a spot cleaning starting from the base."""
         self.robot.start_spot_cleaning()
+
+    def neato_custom_cleaning(self, mode, navigation, category,
+                              zone=None, **kwargs):
+        """Zone cleaning service call."""
+        boundary_id = None
+        if zone is not None:
+            for boundary in self._robot_boundaries['data']['boundaries']:
+                if zone in boundary['name']:
+                    boundary_id = boundary['id']
+            if boundary_id is None:
+                _LOGGER.error(
+                    "Zone '%s' was not found for the robot '%s'",
+                    zone, self._name)
+                return
+
+        self._clean_state = STATE_CLEANING
+        self.robot.start_cleaning(mode, navigation, category, boundary_id)

homeassistant/components/vacuum/services.yaml

@@ -144,3 +144,22 @@ xiaomi_clean_zone:
     repeats:
       description: Number of cleaning repeats for each zone between 1 and 3.
       example: '1'
+
+neato_custom_cleaning:
+  description: Zone Cleaning service call specific to Neato Botvacs.
+  fields:
+    entity_id:
+      description: Name of the vacuum entity. [Required]
+      example: 'vacuum.neato'
+    mode:
+      description: "Set the cleaning mode: 1 for eco and 2 for turbo. Defaults to turbo if not set."
+      example: 2
+    navigation:
+      description: "Set the navigation mode: 1 for normal, 2 for extra care, 3 for deep. Defaults to normal if not set."
+      example: 1
+    category:
+      description: "Whether to use a persistent map or not for cleaning (i.e. No go lines): 2 for no map, 4 for map. Default to using map if not set (and fallback to no map if no map is found)."
+      example: 2
+    zone:
+      description: Only supported on the Botvac D7. Name of the zone to clean. Defaults to no zone i.e. complete house cleanup.
+      example: "Kitchen"