Improve type hints in XiaomiGatewayDevice derived entities (#145605)

homeassistant/components/xiaomi_miio/entity.py

@@ -4,9 +4,10 @@ import datetime
 from enum import Enum
 from functools import partial
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from miio import DeviceException
+from miio.gateway.devices import SubDevice
 
 from homeassistant.const import ATTR_CONNECTIONS, CONF_MAC, CONF_MODEL
 from homeassistant.helpers import device_registry as dr
@@ -18,6 +19,7 @@ from homeassistant.helpers.update_coordinator import (
 )
 
 from .const import ATTR_AVAILABLE, DOMAIN
+from .typing import XiaomiMiioConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -150,10 +152,17 @@ class XiaomiCoordinatedMiioEntity[_T: DataUpdateCoordinator[Any]](
         return time.isoformat()
 
 
-class XiaomiGatewayDevice(CoordinatorEntity, Entity):
+class XiaomiGatewayDevice(
+    CoordinatorEntity[DataUpdateCoordinator[dict[str, bool]]], Entity
+):
     """Representation of a base Xiaomi Gateway Device."""
 
-    def __init__(self, coordinator, sub_device, entry):
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator[dict[str, bool]],
+        sub_device: SubDevice,
+        entry: XiaomiMiioConfigEntry,
+    ) -> None:
         """Initialize the Xiaomi Gateway Device."""
         super().__init__(coordinator)
         self._sub_device = sub_device
@@ -174,6 +183,8 @@ class XiaomiGatewayDevice(CoordinatorEntity, Entity):
     @property
     def device_info(self) -> DeviceInfo:
         """Return the device info of the gateway."""
+        if TYPE_CHECKING:
+            assert self._entry.unique_id is not None
         return DeviceInfo(
             identifiers={(DOMAIN, self._sub_device.sid)},
             via_device=(DOMAIN, self._entry.unique_id),

homeassistant/components/xiaomi_miio/light.py

@@ -18,6 +18,7 @@ from miio import (
     PhilipsEyecare,
     PhilipsMoonlight,
 )
+from miio.gateway.devices.light import LightBulb
 from miio.gateway.gateway import (
     GATEWAY_MODEL_AC_V1,
     GATEWAY_MODEL_AC_V2,
@@ -1093,6 +1094,7 @@ class XiaomiGatewayBulb(XiaomiGatewayDevice, LightEntity):
 
     _attr_color_mode = ColorMode.COLOR_TEMP
     _attr_supported_color_modes = {ColorMode.COLOR_TEMP}
+    _sub_device: LightBulb
 
     @property
     def brightness(self):

homeassistant/components/xiaomi_miio/sensor.py

@@ -8,6 +8,7 @@ import logging
 from typing import TYPE_CHECKING
 
 from miio import AirQualityMonitor, Device as MiioDevice, DeviceException
+from miio.gateway.devices import SubDevice
 from miio.gateway.gateway import (
     GATEWAY_MODEL_AC_V1,
     GATEWAY_MODEL_AC_V2,
@@ -46,6 +47,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from homeassistant.util import dt as dt_util
 
 from . import VacuumCoordinatorDataAttributes
@@ -977,7 +979,13 @@ class XiaomiAirQualityMonitor(XiaomiMiioEntity, SensorEntity):
 class XiaomiGatewaySensor(XiaomiGatewayDevice, SensorEntity):
     """Representation of a XiaomiGatewaySensor."""
 
-    def __init__(self, coordinator, sub_device, entry, description):
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator[dict[str, bool]],
+        sub_device: SubDevice,
+        entry: XiaomiMiioConfigEntry,
+        description: XiaomiMiioSensorDescription,
+    ) -> None:
         """Initialize the XiaomiSensor."""
         super().__init__(coordinator, sub_device, entry)
         self._unique_id = f"{sub_device.sid}-{description.key}"

homeassistant/components/xiaomi_miio/switch.py

@@ -9,6 +9,8 @@ import logging
 from typing import Any
 
 from miio import AirConditioningCompanionV3, ChuangmiPlug, DeviceException, PowerStrip
+from miio.gateway.devices import SubDevice
+from miio.gateway.devices.switch import Switch
 from miio.powerstrip import PowerMode
 import voluptuous as vol
 
@@ -30,6 +32,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, ServiceCall, callback
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import (
     CONF_FLOW_TYPE,
@@ -748,8 +751,15 @@ class XiaomiGatewaySwitch(XiaomiGatewayDevice, SwitchEntity):
     """Representation of a XiaomiGatewaySwitch."""
 
     _attr_device_class = SwitchDeviceClass.SWITCH
+    _sub_device: Switch
 
-    def __init__(self, coordinator, sub_device, entry, variable):
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator[dict[str, bool]],
+        sub_device: SubDevice,
+        entry: XiaomiMiioConfigEntry,
+        variable: str,
+    ) -> None:
         """Initialize the XiaomiSensor."""
         super().__init__(coordinator, sub_device, entry)
         self._channel = GATEWAY_SWITCH_VARS[variable][KEY_CHANNEL]