Merge main

Author: mdeweerd

custom_components/zha_toolkit/default.py

@@ -1,6 +1,12 @@
+from __future__ import annotations
+
 import importlib
 import logging
 import sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from types import ModuleType
 
 LOGGER = logging.getLogger(__name__)
 
@@ -27,13 +33,18 @@ async def default(app, listener, ieee, cmd, data, service, params, event_data):
         module_name = cmd[0]
         cmd = cmd[1]
 
-    LOGGER.debug(
-        f"Trying to import {package_name}.{module_name} to call {cmd}"
-    )
-    m = importlib.import_module(f".{module_name}", package=package_name)
+    def _reload_command_module() -> ModuleType:
+        LOGGER.debug(
+            f"Trying to import {package_name}.{module_name} to call {cmd}"
+        )
+        m = importlib.import_module(f".{module_name}", package=package_name)
 
-    importlib.reload(m)
+        importlib.reload(m)
+        return m
 
+    m = await listener.hass.async_add_import_executor_job(
+        _reload_command_module
+    )
     # Get handler (cmd) in loaded module.
     handler = getattr(m, cmd)
     # Call the handler

custom_components/zha_toolkit/ota.py

@@ -158,28 +158,40 @@ async def download_sonoff_ota(listener, ota_dir):
 
 async def download_zigpy_ota(app, listener):
     LOGGER.debug("Zigpy download procedure starting")
-    for _, (ota, _) in app.ota._listeners.items():
-        if isinstance(ota, zigpy.ota.provider.FileStore):
-            # Skip files provider
-            continue
-        await ota.refresh_firmware_list()
-        for image_key, image in ota._cache.items():
-            url = getattr(image, "url", None)
-            LOGGER.debug("Try getting %r, %r, %r", image_key, url, image)
-            try:
-                img = await app.ota.get_ota_image(
-                    image_key.manufacturer_id, image_key.image_type, model=None
-                )
-                LOGGER.debug("Got image %r", getattr(img, "header", None))
-            except Exception as e:
-                LOGGER.error("%r while getting %r - %s", e, image_key, url)
+    if hasattr(app, "ota") and hasattr(app.ota, "_listeners"):
+        for _, (ota, _) in app.ota._listeners.items():
+            if isinstance(ota, zigpy.ota.provider.FileStore):
+                # Skip files provider
+                continue
+            await ota.refresh_firmware_list()
+            for image_key, image in ota._cache.items():
+                url = getattr(image, "url", None)
+                LOGGER.error("Try getting %r, %r, %r", image_key, url, image)
+                try:
+                    img = await app.ota.get_ota_image(
+                        image_key.manufacturer_id,
+                        image_key.image_type,
+                        model=None,
+                    )
+                    LOGGER.info("Got image %r", getattr(img, "header", None))
+                except Exception as e:
+                    LOGGER.error("%r while getting %r - %s", e, image_key, url)
+    else:
+        LOGGER.warning(
+            "Could not get ota object for download_zigpy_ota, try again"
+        )
 
 
 async def ota_update_images(
     app, listener, ieee, cmd, data, service, params, event_data
 ):
-    for _, (ota, _) in app.ota._listeners.items():
-        await ota.refresh_firmware_list()
+    if hasattr(app, "ota") and hasattr(app.ota, "_listeners"):
+        for _, (ota, _) in app.ota._listeners.items():
+            await ota.refresh_firmware_list()
+    else:
+        LOGGER.warning(
+            "Could not get ota object for ota_update_images, try again"
+        )
 
 
 async def ota_notify(

custom_components/zha_toolkit/utils.py

@@ -34,6 +34,10 @@
     # pylint: disable=ungrouped-imports
     from homeassistant.helpers.json import save_json
 
+if parse_version(HA_VERSION) >= parse_version("2024.6"):
+    from homeassistant.helpers import device_registry as dr
+    from homeassistant.helpers import entity_registry as er
+
 if typing.TYPE_CHECKING:
     VERSION_TIME: float = 0.0
     VERSION: str = "Unknown"
@@ -286,6 +290,8 @@ async def get_ieee(app, listener, ref):
             else get_hass(listener).helpers.entity_registry.async_get(
                 get_hass(listener)
             )
+            if not is_ha_ge("2024.6")
+            else er.async_get(get_hass(listener))
         )
 
         device_registry = (
@@ -297,6 +303,8 @@ async def get_ieee(app, listener, ref):
             else get_hass(listener).helpers.device_registry.async_get(
                 get_hass(listener)
             )
+            if not is_ha_ge("2024.6")
+            else dr.async_get(get_hass(listener))
         )
         registry_device = device_registry.async_get(ref)
 
@@ -356,7 +364,11 @@ def set_state(
     #              entity_id, key, value, stateAttrs)
     if key is not None:
         stateAttrs[key] = value
-        value = None
+        if stateObj is not None:
+            # Copy existing state, to update selected item
+            value = stateObj.state
+        else:
+            value = None
 
     # LOGGER.debug("entity:%s key:%s value:%s attrs:%s",
     #              entity_id, key, value, stateAttrs)