Code refactoring with black style formatter

custom_components/vacuum_zones/__init__.py

@@ -1,28 +1,37 @@
 import homeassistant.helpers.config_validation as cv
 import voluptuous as vol
 from homeassistant.const import CONF_ENTITY_ID, CONF_SEQUENCE
-from homeassistant.helpers.typing import HomeAssistantType
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.discovery import async_load_platform
 
-DOMAIN = 'vacuum_zones'
+DOMAIN = "vacuum_zones"
+PLATFORMS = ["sensor"]
 
-CONFIG_SCHEMA = vol.Schema({
-    DOMAIN: vol.Schema({
-        vol.Required(CONF_ENTITY_ID): cv.entity_id,
-        vol.Required('zones'): {
-            cv.string: vol.Schema({
-                vol.Optional('room'): vol.Any(list, int),
-                vol.Optional('zone'): list,
-                vol.Optional('repeats'): int,
-                vol.Optional('goto'): list,
-                vol.Optional(CONF_SEQUENCE): cv.SCRIPT_SCHEMA
-            })
-        }
-    })
-}, extra=vol.ALLOW_EXTRA)
+CONFIG_SCHEMA = vol.Schema(
+    {
+        DOMAIN: vol.Schema(
+            {
+                vol.Required(CONF_ENTITY_ID): cv.entity_id,
+                vol.Required("zones"): {
+                    cv.string: vol.Schema(
+                        {
+                            vol.Optional("room"): vol.Any(list, int),
+                            vol.Optional("zone"): list,
+                            vol.Optional("repeats"): int,
+                            vol.Optional("goto"): list,
+                            vol.Optional(CONF_SEQUENCE): cv.SCRIPT_SCHEMA,
+                        }
+                    )
+                },
+            }
+        )
+    },
+    extra=vol.ALLOW_EXTRA,
+)
 
 
-async def async_setup(hass: HomeAssistantType, hass_config: dict):
-    hass.async_create_task(hass.helpers.discovery.async_load_platform(
-        'vacuum', DOMAIN, hass_config[DOMAIN], hass_config
-    ))
+async def async_setup(hass: HomeAssistant, config: dict):
+    hass.async_create_task(
+        async_load_platform("vacuum", DOMAIN, config[DOMAIN], config)
+    )
     return True

custom_components/vacuum_zones/vacuum.py

@@ -1,29 +1,39 @@
 from typing import List
 
-from homeassistant.components.vacuum import StateVacuumEntity, SUPPORT_START, \
-    SUPPORT_STOP, DOMAIN as VACUUM_DOMAIN, STATE_CLEANING, STATE_RETURNING, \
-    STATE_DOCKED
-from homeassistant.const import CONF_SEQUENCE, STATE_IDLE, STATE_PAUSED, \
-    EVENT_STATE_CHANGED, ATTR_ENTITY_ID
+from homeassistant.components.vacuum import (
+    StateVacuumEntity,
+    SUPPORT_START,
+    SUPPORT_STOP,
+    DOMAIN as VACUUM_DOMAIN,
+    STATE_CLEANING,
+    STATE_RETURNING,
+    STATE_DOCKED,
+)
+from homeassistant.const import (
+    CONF_SEQUENCE,
+    STATE_IDLE,
+    STATE_PAUSED,
+    EVENT_STATE_CHANGED,
+    ATTR_ENTITY_ID,
+)
 from homeassistant.core import Event, State
 from homeassistant.helpers.script import Script
 
 
-async def async_setup_platform(hass, _, async_add_entities,
-                               discovery_info=None):
-    entity_id = discovery_info['entity_id']
+async def async_setup_platform(hass, _, async_add_entities, discovery_info=None):
+    entity_id = discovery_info["entity_id"]
     queue: List[XiaomiVacuum] = []
     entities = [
         XiaomiVacuum(name, config, entity_id, queue)
-        for name, config in discovery_info['zones'].items()
+        for name, config in discovery_info["zones"].items()
     ]
     async_add_entities(entities)
 
     async def state_changed_event_listener(event: Event):
         if entity_id != event.data.get(ATTR_ENTITY_ID) or not queue:
             return
 
-        new_state: State = event.data.get('new_state')
+        new_state: State = event.data.get("new_state")
         if new_state.state not in (STATE_RETURNING, STATE_DOCKED):
             return
 
@@ -40,73 +50,76 @@ async def state_changed_event_listener(event: Event):
 
 
 class XiaomiVacuum(StateVacuumEntity):
-    _state = STATE_IDLE
+    _attr_state = STATE_IDLE
+    _attr_supported_features = SUPPORT_START | SUPPORT_STOP
+
     script = None
 
     def __init__(self, name: str, config: dict, entity_id: str, queue: list):
-        config['name'] = name
-        config['entity_id'] = entity_id
+        self._attr_name = name
+        config["entity_id"] = entity_id
         self.config = config
         self.queue = queue
 
     async def async_added_to_hass(self):
         if CONF_SEQUENCE in self.config:
-            self.script = Script(self.hass, self.config[CONF_SEQUENCE],
-                                 self.name, VACUUM_DOMAIN)
-
-    @property
-    def name(self):
-        return self.config['name']
-
-    @property
-    def state(self):
-        return self._state
-
-    @property
-    def supported_features(self):
-        return SUPPORT_START | SUPPORT_STOP
+            self.script = Script(
+                self.hass, self.config[CONF_SEQUENCE], self.name, VACUUM_DOMAIN
+            )
 
     async def internal_start(self):
-        self._state = STATE_CLEANING
+        self._attr_state = STATE_CLEANING
         await self.async_update_ha_state()
 
         if self.script:
             await self.script.async_run()
 
-        if 'room' in self.config:
+        if "room" in self.config:
             await self.hass.services.async_call(
-                'xiaomi_miio', 'vacuum_clean_segment', {
-                    'entity_id': self.config['entity_id'],
-                    'segments': self.config['room'],
-                }, blocking=True)
-
-        elif 'zone' in self.config:
+                "xiaomi_miio",
+                "vacuum_clean_segment",
+                {
+                    "entity_id": self.config["entity_id"],
+                    "segments": self.config["room"],
+                },
+                blocking=True,
+            )
+
+        elif "zone" in self.config:
             await self.hass.services.async_call(
-                'xiaomi_miio', 'vacuum_clean_zone', {
-                    'entity_id': self.config['entity_id'],
-                    'zone': self.config['zone'],
-                    'repeats': self.config.get('repeats', 1)
-                }, blocking=True)
-
-        elif 'goto' in self.config:
+                "xiaomi_miio",
+                "vacuum_clean_zone",
+                {
+                    "entity_id": self.config["entity_id"],
+                    "zone": self.config["zone"],
+                    "repeats": self.config.get("repeats", 1),
+                },
+                blocking=True,
+            )
+
+        elif "goto" in self.config:
             await self.hass.services.async_call(
-                'xiaomi_miio', 'vacuum_goto', {
-                    'entity_id': self.config['entity_id'],
-                    'x_coord': self.config['goto'][0],
-                    'y_coord': self.config['goto'][1],
-                }, blocking=True)
+                "xiaomi_miio",
+                "vacuum_goto",
+                {
+                    "entity_id": self.config["entity_id"],
+                    "x_coord": self.config["goto"][0],
+                    "y_coord": self.config["goto"][1],
+                },
+                blocking=True,
+            )
 
     async def internal_stop(self):
-        self._state = STATE_IDLE
-        await self.async_update_ha_state()
+        self._attr_state = STATE_IDLE
+        self.async_write_ha_state()
 
     async def async_start(self):
         self.queue.append(self)
 
-        state = self.hass.states.get(self.config['entity_id'])
+        state = self.hass.states.get(self.config["entity_id"])
         if len(self.queue) > 1 or state == STATE_CLEANING:
-            self._state = STATE_PAUSED
-            await self.async_update_ha_state()
+            self._attr_state = STATE_PAUSED
+            self.async_write_ha_state()
             return
 
         await self.internal_start()