From 285c0d2c024e13fa1f3add8808572bcc95ac0cb0 Mon Sep 17 00:00:00 2001 From: mdeweerd Date: Thu, 8 Aug 2024 19:03:13 +0200 Subject: [PATCH] Get gateway using helper function --- custom_components/zha_toolkit/__init__.py | 5 +---- custom_components/zha_toolkit/utils.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/custom_components/zha_toolkit/__init__.py b/custom_components/zha_toolkit/__init__.py index ebc8b42..5a4057d 100644 --- a/custom_components/zha_toolkit/__init__.py +++ b/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( diff --git a/custom_components/zha_toolkit/utils.py b/custom_components/zha_toolkit/utils.py index 39e3e4f..3a471db 100644 --- a/custom_components/zha_toolkit/utils.py +++ b/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