Add zha typing [core.gateway] (2) (#68644)

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:

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
 

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(),