Get gateway using helper function

Author: mdeweerd

custom_components/zha_toolkit/__init__.py

@@ -678,10 +678,7 @@ async def toolkit_service(service):
 
         zha = hass_ref.data["zha"]
         zha_gw: Optional[ZHAGateway] = None
-        if isinstance(zha, dict):
-            zha_gw = zha.get("zha_gateway", None)
-        else:
-            zha_gw = zha.gateway
+        zha_gw = u.get_zha_gateway(hass)
 
         if zha_gw is None:
             LOGGER.error(

custom_components/zha_toolkit/utils.py

@@ -11,12 +11,15 @@
 
 import aiofiles
 import zigpy
+from homeassistant.core import HomeAssistant
 
 try:
     from homeassistant.components.zha import Gateway as ZHAGateway
 except ImportError:
     from homeassistant.components.zha.core.gateway import ZHAGateway
 
+from homeassistant.components import zha
+from homeassistant.components.zha import helpers as zha_helpers
 from homeassistant.util import dt as dt_util
 from pkg_resources import get_distribution, parse_version
 from zigpy import types as t
@@ -50,6 +53,15 @@
     MANIFEST: dict[str, str | list[str]] = {}
 
 
+def get_zha_gateway(hass: HomeAssistant) -> ZHAGateway:
+    """Get the ZHA gateway object."""
+    if parse_version(HA_VERSION) >= parse_version("2024.8"):
+        return zha_helpers.get_zha_gateway(hass)
+    if isinstance(zha, dict):
+        return zha.get("zha_gateway", None)
+    return zha.gateway
+
+
 def getHaVersion() -> str:
     """Get HA Version"""
     return HA_VERSION