attempt at cleaning up device removal

anotherjulien · Mar 10, 2024 · 2da12d0 · 2da12d0
1 parent 936cb93
commit 2da12d0
Showing 1 changed file with 86 additions and 25 deletions.
diff --git a/custom_components/myhome/__init__.py b/custom_components/myhome/__init__.py
@@ -1,4 +1,5 @@
 """ MyHOME integration. """
+
 import aiofiles
 import yaml
 
@@ -45,8 +46,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     if entry.data[CONF_MAC] not in hass.data[DOMAIN]:
         hass.data[DOMAIN][entry.data[CONF_MAC]] = {}
 
-    _config_file_path = str(entry.options[CONF_FILE_PATH]) if CONF_FILE_PATH in entry.options else "/config/myhome.yaml"
-    _generate_events = entry.options[CONF_GENERATE_EVENTS] if CONF_GENERATE_EVENTS in entry.options else False
+    _config_file_path = (
+        str(entry.options[CONF_FILE_PATH])
+        if CONF_FILE_PATH in entry.options
+        else "/config/myhome.yaml"
+    )
+    _generate_events = (
+        entry.options[CONF_GENERATE_EVENTS]
+        if CONF_GENERATE_EVENTS in entry.options
+        else False
+    )
 
     try:
         async with aiofiles.open(_config_file_path, mode="r") as yaml_file:
@@ -56,26 +65,39 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
         return False
 
     if entry.data[CONF_MAC] in _validated_config:
-        hass.data[DOMAIN][entry.data[CONF_MAC]] = _validated_config[entry.data[CONF_MAC]]
+        hass.data[DOMAIN][entry.data[CONF_MAC]] = _validated_config[
+            entry.data[CONF_MAC]
+        ]
     else:
         return False
 
     # Migrating the config entry's unique_id if it was not formated to the recommended hass standard
     if entry.unique_id != dr.format_mac(entry.unique_id):
-        hass.config_entries.async_update_entry(entry, unique_id=dr.format_mac(entry.unique_id))
+        hass.config_entries.async_update_entry(
+            entry, unique_id=dr.format_mac(entry.unique_id)
+        )
         LOGGER.warning("Migrating config entry unique_id to %s", entry.unique_id)
 
-    hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY] = MyHOMEGatewayHandler(hass=hass, config_entry=entry, generate_events=_generate_events)
+    hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY] = MyHOMEGatewayHandler(
+        hass=hass, config_entry=entry, generate_events=_generate_events
+    )
 
     try:
-        tests_results = await hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].test()
+        tests_results = await hass.data[DOMAIN][entry.data[CONF_MAC]][
+            CONF_ENTITY
+        ].test()
     except OSError as ose:
         _gateway_handler = hass.data[DOMAIN].pop(CONF_GATEWAY)
         _host = _gateway_handler.gateway.host
-        raise ConfigEntryNotReady(f"Gateway cannot be reached at {_host}, make sure its address is correct.") from ose
+        raise ConfigEntryNotReady(
+            f"Gateway cannot be reached at {_host}, make sure its address is correct."
+        ) from ose
 
     if not tests_results["Success"]:
-        if tests_results["Message"] == "password_error" or tests_results["Message"] == "password_required":
+        if (
+            tests_results["Message"] == "password_error"
+            or tests_results["Message"] == "password_required"
+        ):
             hass.async_create_task(
                 hass.config_entries.flow.async_init(
                     DOMAIN,
@@ -86,47 +108,70 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
         del hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY]
         return False
 
-    _command_worker_count = int(entry.options[CONF_WORKER_COUNT]) if CONF_WORKER_COUNT in entry.options else 1
+    _command_worker_count = (
+        int(entry.options[CONF_WORKER_COUNT])
+        if CONF_WORKER_COUNT in entry.options
+        else 1
+    )
 
     entity_registry = er.async_get(hass)
     device_registry = dr.async_get(hass)
 
-    gateway_entry = device_registry.async_get_or_create(
+    gateway_device_entry = device_registry.async_get_or_create(
         config_entry_id=entry.entry_id,
         connections={(dr.CONNECTION_NETWORK_MAC, entry.data[CONF_MAC])},
-        identifiers={(DOMAIN, hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].unique_id)},
+        identifiers={
+            (DOMAIN, hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].unique_id)
+        },
         manufacturer=hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].manufacturer,
         name=hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].name,
         model=hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].model,
         sw_version=hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].firmware,
     )
 
-    await hass.config_entries.async_forward_entry_setups(entry, hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_PLATFORMS].keys())
+    await hass.config_entries.async_forward_entry_setups(
+        entry, hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_PLATFORMS].keys()
+    )
 
-    hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].listening_worker = hass.loop.create_task(hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].listening_loop())
+    hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].listening_worker = (
+        hass.loop.create_task(
+            hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].listening_loop()
+        )
+    )
     for i in range(_command_worker_count):
-        hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].sending_workers.append(hass.loop.create_task(hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].sending_loop(i)))
+        hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].sending_workers.append(
+            hass.loop.create_task(
+                hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_ENTITY].sending_loop(i)
+            )
+        )
 
     # Pruning lose entities and devices from the registry
     entity_entries = er.async_entries_for_config_entry(entity_registry, entry.entry_id)
 
     entities_to_be_removed = []
-    devices_to_be_removed = [device_entry.id for device_entry in device_registry.devices.values() if entry.entry_id in device_entry.config_entries]
-
-    if gateway_entry.id in devices_to_be_removed:
-        devices_to_be_removed.remove(gateway_entry.id)
+    devices_to_be_removed = [
+        device_entry.id
+        for device_entry in device_registry.devices.values()
+        if entry.entry_id in device_entry.config_entries
+    ]
 
     configured_entities = []
 
     for _platform in hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_PLATFORMS].keys():
-        for _device in hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_PLATFORMS][_platform].keys():
-            for _entity_name in hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_PLATFORMS][_platform][_device][CONF_ENTITIES]:
+        for _device in hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_PLATFORMS][
+            _platform
+        ].keys():
+            for _entity_name in hass.data[DOMAIN][entry.data[CONF_MAC]][CONF_PLATFORMS][
+                _platform
+            ][_device][CONF_ENTITIES]:
                 if _entity_name != _platform:
                     configured_entities.append(
                         f"{entry.data[CONF_MAC]}-{_device}-{_entity_name}"
                     )  # extrapolating _attr_unique_id out of the entity's place in the config data structure
                 else:
-                    configured_entities.append(f"{entry.data[CONF_MAC]}-{_device}")  # extrapolating _attr_unique_id out of the entity's place in the config data structure
+                    configured_entities.append(
+                        f"{entry.data[CONF_MAC]}-{_device}"
+                    )  # extrapolating _attr_unique_id out of the entity's place in the config data structure
 
     for entity_entry in entity_entries:
         if entity_entry.unique_id in configured_entities:
@@ -138,8 +183,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     for enity_id in entities_to_be_removed:
         entity_registry.async_remove(enity_id)
 
+    if gateway_device_entry.id in devices_to_be_removed:
+        devices_to_be_removed.remove(gateway_device_entry.id)
+
     for device_id in devices_to_be_removed:
-        if len(er.async_entries_for_device(entity_registry, device_id, include_disabled_entities=True)) == 0:
+        if (
+            len(
+                er.async_entries_for_device(
+                    entity_registry, device_id, include_disabled_entities=True
+                )
+            )
+            == 0
+        ):
             device_registry.async_remove_device(device_id)
 
     # Defining the services
@@ -159,7 +214,9 @@ async def handle_sync_time(call):
                 gateway = mac
         timezone = hass.config.as_dict()["time_zone"]
         if gateway in hass.data[DOMAIN]:
-            await hass.data[DOMAIN][gateway][CONF_ENTITY].send(OWNGatewayCommand.set_datetime_to_now(timezone))
+            await hass.data[DOMAIN][gateway][CONF_ENTITY].send(
+                OWNGatewayCommand.set_datetime_to_now(timezone)
+            )
         else:
             LOGGER.error(
                 "Gateway `%s` not found, could not send time synchronisation message.",
@@ -198,10 +255,14 @@ async def handle_send_message(call):
                         )
                         await hass.data[DOMAIN][gateway][CONF_ENTITY].send(own_message)
                 else:
-                    LOGGER.error("Could not parse message `%s`, not sending it.", message)
+                    LOGGER.error(
+                        "Could not parse message `%s`, not sending it.", message
+                    )
                     return False
         else:
-            LOGGER.error("Gateway `%s` not found, could not send message `%s`.", gateway, message)
+            LOGGER.error(
+                "Gateway `%s` not found, could not send message `%s`.", gateway, message
+            )
             return False
 
     hass.services.async_register(DOMAIN, "send_message", handle_send_message)