diff --git a/homeassistant/components/zha/__init__.py b/homeassistant/components/zha/__init__.py index 4f1e80e0a7b..c512e104ae8 100644 --- a/homeassistant/components/zha/__init__.py +++ b/homeassistant/components/zha/__init__.py @@ -109,8 +109,8 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b device_registry = await hass.helpers.device_registry.async_get_registry() device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(CONNECTION_ZIGBEE, str(zha_gateway.application_controller.ieee))}, # type: ignore[attr-defined] - identifiers={(DOMAIN, str(zha_gateway.application_controller.ieee))}, # type: ignore[attr-defined] + connections={(CONNECTION_ZIGBEE, str(zha_gateway.application_controller.ieee))}, + identifiers={(DOMAIN, str(zha_gateway.application_controller.ieee))}, name="Zigbee Coordinator", manufacturer="ZHA", model=zha_gateway.radio_description, @@ -120,8 +120,9 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b async def async_zha_shutdown(event): """Handle shutdown tasks.""" - await zha_data[DATA_ZHA_GATEWAY].shutdown() - await zha_data[DATA_ZHA_GATEWAY].async_update_device_storage() + zha_gateway: ZHAGateway = zha_data[DATA_ZHA_GATEWAY] + await zha_gateway.shutdown() + await zha_gateway.async_update_device_storage() zha_data[DATA_ZHA_SHUTDOWN_TASK] = hass.bus.async_listen_once( ha_const.EVENT_HOMEASSISTANT_STOP, async_zha_shutdown @@ -132,8 +133,9 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: """Unload ZHA config entry.""" - await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].shutdown() - await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_update_device_storage() + zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] + await zha_gateway.shutdown() + await zha_gateway.async_update_device_storage() GROUP_PROBE.cleanup() api.async_unload_api(hass) @@ -153,7 +155,8 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> async def async_load_entities(hass: HomeAssistant) -> None: """Load entities after integration was setup.""" - await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_initialize_devices_and_entities() + zha_gateway: ZHAGateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] + await zha_gateway.async_initialize_devices_and_entities() to_setup = hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED] results = await asyncio.gather(*to_setup, return_exceptions=True) for res in results: diff --git a/homeassistant/components/zha/core/const.py b/homeassistant/components/zha/core/const.py index 9216b6dd600..b3f06b9eba6 100644 --- a/homeassistant/components/zha/core/const.py +++ b/homeassistant/components/zha/core/const.py @@ -6,6 +6,7 @@ import logging import bellows.zigbee.application import voluptuous as vol +import zigpy.application from zigpy.config import CONF_DEVICE_PATH # noqa: F401 # pylint: disable=unused-import import zigpy.types as t import zigpy_deconz.zigbee.application @@ -16,8 +17,6 @@ import zigpy_znp.zigbee.application from homeassistant.const import Platform import homeassistant.helpers.config_validation as cv -from .typing import CALLABLE_T - ATTR_ARGS = "args" ATTR_ATTRIBUTE = "attribute" ATTR_ATTRIBUTE_ID = "attribute_id" @@ -224,6 +223,8 @@ ZHA_CONFIG_SCHEMAS = { ZHA_ALARM_OPTIONS: CONF_ZHA_ALARM_SCHEMA, } +_ControllerClsType = type[zigpy.application.ControllerApplication] + class RadioType(enum.Enum): """Possible options for radio type.""" @@ -262,13 +263,13 @@ class RadioType(enum.Enum): return radio.name raise ValueError - def __init__(self, description: str, controller_cls: CALLABLE_T) -> None: + def __init__(self, description: str, controller_cls: _ControllerClsType) -> None: """Init instance.""" self._desc = description self._ctrl_cls = controller_cls @property - def controller(self) -> CALLABLE_T: + def controller(self) -> _ControllerClsType: """Return controller class.""" return self._ctrl_cls diff --git a/homeassistant/components/zha/core/gateway.py b/homeassistant/components/zha/core/gateway.py index 636e161d45c..6c600bf93d6 100644 --- a/homeassistant/components/zha/core/gateway.py +++ b/homeassistant/components/zha/core/gateway.py @@ -14,6 +14,7 @@ import traceback from typing import TYPE_CHECKING, Any, Union from serial import SerialException +from zigpy.application import ControllerApplication from zigpy.config import CONF_DEVICE import zigpy.device import zigpy.endpoint @@ -26,10 +27,12 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers.device_registry import ( CONNECTION_ZIGBEE, + DeviceRegistry, async_get_registry as get_dev_reg, ) from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.entity_registry import ( + EntityRegistry, async_entries_for_device, async_get_registry as get_ent_reg, ) @@ -93,6 +96,7 @@ if TYPE_CHECKING: from logging import Filter, LogRecord from ..entity import ZhaEntity + from .store import ZhaStorage _LogFilterType = Union[Filter, Callable[[LogRecord], int]] @@ -116,6 +120,13 @@ class DevicePairingStatus(Enum): class ZHAGateway: """Gateway that handles events that happen on the ZHA Zigbee network.""" + # -- Set in async_initialize -- + zha_storage: ZhaStorage + ha_device_registry: DeviceRegistry + ha_entity_registry: EntityRegistry + application_controller: ControllerApplication + radio_description: str + def __init__( self, hass: HomeAssistant, config: ConfigType, config_entry: ConfigEntry ) -> None: @@ -128,11 +139,6 @@ class ZHAGateway: self._device_registry: collections.defaultdict[ EUI64, list[EntityReference] ] = collections.defaultdict(list) - self.zha_storage = None - self.ha_device_registry = None - self.ha_entity_registry = None - self.application_controller = None - self.radio_description = None self._log_levels: dict[str, dict[str, int]] = { DEBUG_LEVEL_ORIGINAL: async_capture_log_levels(), DEBUG_LEVEL_CURRENT: async_capture_log_levels(),