Improve CoordinatorEntity typing (#68441)

homeassistant/components/co2signal/sensor.py

@@ -9,13 +9,13 @@ from homeassistant.components.sensor import SensorEntity, SensorStateClass
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_ATTRIBUTION, PERCENTAGE
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import update_coordinator
 from homeassistant.helpers.device_registry import DeviceEntryType
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import CO2SignalCoordinator, CO2SignalResponse
+from . import CO2SignalCoordinator
 from .const import ATTRIBUTION, DOMAIN
 
 SCAN_INTERVAL = timedelta(minutes=3)
@@ -55,7 +55,7 @@ async def async_setup_entry(
     async_add_entities(CO2Sensor(coordinator, description) for description in SENSORS)
 
 
-class CO2Sensor(update_coordinator.CoordinatorEntity[CO2SignalResponse], SensorEntity):
+class CO2Sensor(CoordinatorEntity[CO2SignalCoordinator], SensorEntity):
     """Implementation of the CO2Signal sensor."""
 
     _attr_state_class = SensorStateClass.MEASUREMENT

homeassistant/components/fjaraskupan/binary_sensor.py

@@ -67,7 +67,7 @@ async def async_setup_entry(
     async_setup_entry_platform(hass, config_entry, async_add_entities, _constructor)
 
 
-class BinarySensor(CoordinatorEntity[State], BinarySensorEntity):
+class BinarySensor(CoordinatorEntity[DataUpdateCoordinator[State]], BinarySensorEntity):
     """Grease filter sensor."""
 
     entity_description: EntityDescription

homeassistant/components/fjaraskupan/fan.py

@@ -70,7 +70,7 @@ async def async_setup_entry(
     async_setup_entry_platform(hass, config_entry, async_add_entities, _constructor)
 
 
-class Fan(CoordinatorEntity[State], FanEntity):
+class Fan(CoordinatorEntity[DataUpdateCoordinator[State]], FanEntity):
     """Fan entity."""
 
     def __init__(

homeassistant/components/fjaraskupan/light.py

@@ -37,7 +37,7 @@ async def async_setup_entry(
     async_setup_entry_platform(hass, config_entry, async_add_entities, _constructor)
 
 
-class Light(CoordinatorEntity[State], LightEntity):
+class Light(CoordinatorEntity[DataUpdateCoordinator[State]], LightEntity):
     """Light device."""
 
     def __init__(

homeassistant/components/fjaraskupan/number.py

@@ -34,7 +34,9 @@ async def async_setup_entry(
     async_setup_entry_platform(hass, config_entry, async_add_entities, _constructor)
 
 
-class PeriodicVentingTime(CoordinatorEntity[State], NumberEntity):
+class PeriodicVentingTime(
+    CoordinatorEntity[DataUpdateCoordinator[State]], NumberEntity
+):
     """Periodic Venting."""
 
     _attr_max_value: float = 59

homeassistant/components/fjaraskupan/sensor.py

@@ -39,7 +39,7 @@ async def async_setup_entry(
     async_setup_entry_platform(hass, config_entry, async_add_entities, _constructor)
 
 
-class RssiSensor(CoordinatorEntity[State], SensorEntity):
+class RssiSensor(CoordinatorEntity[DataUpdateCoordinator[State]], SensorEntity):
     """Sensor device."""
 
     def __init__(

homeassistant/components/github/sensor.py

@@ -172,12 +172,11 @@ async def async_setup_entry(
     )
 
 
-class GitHubSensorEntity(CoordinatorEntity[dict[str, Any]], SensorEntity):
+class GitHubSensorEntity(CoordinatorEntity[GitHubDataUpdateCoordinator], SensorEntity):
     """Defines a GitHub sensor entity."""
 
     _attr_attribution = "Data provided by the GitHub API"
 
-    coordinator: GitHubDataUpdateCoordinator
     entity_description: GitHubSensorEntityDescription
 
     def __init__(

homeassistant/components/hassio/entity.py

@@ -11,7 +11,7 @@ from . import DOMAIN, HassioDataUpdateCoordinator
 from .const import ATTR_SLUG, DATA_KEY_ADDONS, DATA_KEY_OS
 
 
-class HassioAddonEntity(CoordinatorEntity):
+class HassioAddonEntity(CoordinatorEntity[HassioDataUpdateCoordinator]):
     """Base entity for a Hass.io add-on."""
 
     def __init__(
@@ -38,7 +38,7 @@ class HassioAddonEntity(CoordinatorEntity):
         )
 
 
-class HassioOSEntity(CoordinatorEntity):
+class HassioOSEntity(CoordinatorEntity[HassioDataUpdateCoordinator]):
     """Base Entity for Hass.io OS."""
 
     def __init__(

homeassistant/components/homewizard/sensor.py

@@ -142,7 +142,7 @@ async def async_setup_entry(
     async_add_entities(entities)
 
 
-class HWEnergySensor(CoordinatorEntity[DeviceResponseEntry], SensorEntity):
+class HWEnergySensor(CoordinatorEntity[HWEnergyDeviceUpdateCoordinator], SensorEntity):
     """Representation of a HomeWizard Sensor."""
 
     def __init__(

homeassistant/components/homewizard/switch.py

@@ -31,11 +31,11 @@ async def async_setup_entry(
         )
 
 
-class HWEnergySwitchEntity(CoordinatorEntity, SwitchEntity):
+class HWEnergySwitchEntity(
+    CoordinatorEntity[HWEnergyDeviceUpdateCoordinator], SwitchEntity
+):
     """Representation switchable entity."""
 
-    coordinator: HWEnergyDeviceUpdateCoordinator
-
     def __init__(
         self,
         coordinator: HWEnergyDeviceUpdateCoordinator,

homeassistant/components/kraken/sensor.py

@@ -11,7 +11,10 @@ from homeassistant.helpers import device_registry
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
 
 from . import KrakenData
 from .const import (
@@ -86,7 +89,9 @@ async def async_setup_entry(
     )
 
 
-class KrakenSensor(CoordinatorEntity[Optional[KrakenResponse]], SensorEntity):
+class KrakenSensor(
+    CoordinatorEntity[DataUpdateCoordinator[Optional[KrakenResponse]]], SensorEntity
+):
     """Define a Kraken sensor."""
 
     entity_description: KrakenSensorEntityDescription

homeassistant/components/met/weather.py

@@ -45,7 +45,7 @@ from homeassistant.util.distance import convert as convert_distance
 from homeassistant.util.pressure import convert as convert_pressure
 from homeassistant.util.speed import convert as convert_speed
 
-from . import MetDataUpdateCoordinator, MetWeatherData
+from . import MetDataUpdateCoordinator
 from .const import (
     ATTR_FORECAST_PRECIPITATION,
     ATTR_MAP,
@@ -127,11 +127,9 @@ def format_condition(condition: str) -> str:
     return condition
 
 
-class MetWeather(CoordinatorEntity[MetWeatherData], WeatherEntity):
+class MetWeather(CoordinatorEntity[MetDataUpdateCoordinator], WeatherEntity):
     """Implementation of a Met.no weather condition."""
 
-    coordinator: MetDataUpdateCoordinator
-
     def __init__(
         self,
         coordinator: MetDataUpdateCoordinator,

homeassistant/components/modern_forms/__init__.py

@@ -126,8 +126,6 @@ class ModernFormsDataUpdateCoordinator(DataUpdateCoordinator[ModernFormsDeviceSt
 class ModernFormsDeviceEntity(CoordinatorEntity[ModernFormsDataUpdateCoordinator]):
     """Defines a Modern Forms device entity."""
 
-    coordinator: ModernFormsDataUpdateCoordinator
-
     def __init__(
         self,
         *,

homeassistant/components/plugwise/entity.py

@@ -12,14 +12,12 @@ from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
-from .coordinator import PlugwiseData, PlugwiseDataUpdateCoordinator
+from .coordinator import PlugwiseDataUpdateCoordinator
 
 
-class PlugwiseEntity(CoordinatorEntity[PlugwiseData]):
+class PlugwiseEntity(CoordinatorEntity[PlugwiseDataUpdateCoordinator]):
     """Represent a PlugWise Entity."""
 
-    coordinator: PlugwiseDataUpdateCoordinator
-
     def __init__(
         self,
         coordinator: PlugwiseDataUpdateCoordinator,

homeassistant/components/powerwall/entity.py

@@ -1,7 +1,10 @@
 """The Tesla Powerwall integration base entity."""
 
 from homeassistant.helpers.entity import DeviceInfo
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
 
 from .const import (
     DOMAIN,
@@ -13,7 +16,7 @@ from .const import (
 from .models import PowerwallData, PowerwallRuntimeData
 
 
-class PowerWallEntity(CoordinatorEntity[PowerwallData]):
+class PowerWallEntity(CoordinatorEntity[DataUpdateCoordinator[PowerwallData]]):
     """Base class for powerwall entities."""
 
     def __init__(self, powerwall_data: PowerwallRuntimeData) -> None:

homeassistant/components/pure_energie/sensor.py

@@ -78,10 +78,11 @@ async def async_setup_entry(
     )
 
 
-class PureEnergieSensorEntity(CoordinatorEntity[PureEnergieData], SensorEntity):
+class PureEnergieSensorEntity(
+    CoordinatorEntity[PureEnergieDataUpdateCoordinator], SensorEntity
+):
     """Defines an Pure Energie sensor."""
 
-    coordinator: PureEnergieDataUpdateCoordinator
     entity_description: PureEnergieSensorEntityDescription
 
     def __init__(

homeassistant/components/renault/renault_coordinator.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 from collections.abc import Awaitable, Callable
 from datetime import timedelta
 import logging
-from typing import TypeVar
+from typing import Optional, TypeVar
 
 from renault_api.kamereon.exceptions import (
     AccessDeniedException,
@@ -16,7 +16,7 @@ from renault_api.kamereon.models import KamereonVehicleDataAttributes
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-T = TypeVar("T", bound=KamereonVehicleDataAttributes)
+T = TypeVar("T", bound=Optional[KamereonVehicleDataAttributes])
 
 
 class RenaultDataUpdateCoordinator(DataUpdateCoordinator[T]):

homeassistant/components/renault/renault_entities.py

@@ -2,14 +2,14 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Optional, cast
+from typing import cast
 
 from homeassistant.const import ATTR_NAME
 from homeassistant.helpers.entity import Entity, EntityDescription
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .renault_coordinator import T
+from .renault_coordinator import RenaultDataUpdateCoordinator, T
 from .renault_vehicle import RenaultVehicleProxy
 
 
@@ -50,7 +50,9 @@ class RenaultEntity(Entity):
         return f"{self.vehicle.device_info[ATTR_NAME]} {self.entity_description.name}"
 
 
-class RenaultDataEntity(CoordinatorEntity[Optional[T]], RenaultEntity):
+class RenaultDataEntity(
+    CoordinatorEntity[RenaultDataUpdateCoordinator[T]], RenaultEntity
+):
     """Implementation of a Renault entity with a data coordinator."""
 
     def __init__(
@@ -65,5 +67,5 @@ class RenaultDataEntity(CoordinatorEntity[Optional[T]], RenaultEntity):
     def _get_data_attr(self, key: str) -> StateType:
         """Return the attribute value from the coordinator data."""
         if self.coordinator.data is None:
-            return None
+            return None  # type: ignore[unreachable]
         return cast(StateType, getattr(self.coordinator.data, key))

homeassistant/components/renault/sensor.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, Generic, cast
 
 from renault_api.kamereon.enums import ChargeState, PlugState
 from renault_api.kamereon.models import (
@@ -45,18 +45,18 @@ from .renault_vehicle import RenaultVehicleProxy
 
 
 @dataclass
-class RenaultSensorRequiredKeysMixin:
+class RenaultSensorRequiredKeysMixin(Generic[T]):
     """Mixin for required keys."""
 
     data_key: str
-    entity_class: type[RenaultSensor]
+    entity_class: type[RenaultSensor[T]]
 
 
 @dataclass
 class RenaultSensorEntityDescription(
     SensorEntityDescription,
     RenaultDataEntityDescription,
-    RenaultSensorRequiredKeysMixin,
+    RenaultSensorRequiredKeysMixin[T],
 ):
     """Class describing Renault sensor entities."""
 
@@ -73,7 +73,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Renault entities from config entry."""
     proxy: RenaultHub = hass.data[DOMAIN][config_entry.entry_id]
-    entities: list[RenaultSensor] = [
+    entities: list[RenaultSensor[Any]] = [
         description.entity_class(vehicle, description)
         for vehicle in proxy.vehicles.values()
         for description in SENSOR_TYPES
@@ -87,7 +87,7 @@ async def async_setup_entry(
 class RenaultSensor(RenaultDataEntity[T], SensorEntity):
     """Mixin for sensor specific attributes."""
 
-    entity_description: RenaultSensorEntityDescription
+    entity_description: RenaultSensorEntityDescription[T]
 
     @property
     def data(self) -> StateType:
@@ -157,7 +157,7 @@ def _get_utc_value(entity: RenaultSensor[T]) -> datetime:
     return as_utc(original_dt)
 
 
-SENSOR_TYPES: tuple[RenaultSensorEntityDescription, ...] = (
+SENSOR_TYPES: tuple[RenaultSensorEntityDescription[Any], ...] = (
     RenaultSensorEntityDescription(
         key="battery_level",
         coordinator="battery",

homeassistant/components/template/trigger_entity.py

@@ -13,13 +13,14 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import TemplateError
-from homeassistant.helpers import template, update_coordinator
+from homeassistant.helpers import template
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import TriggerUpdateCoordinator
 from .const import CONF_ATTRIBUTES, CONF_AVAILABILITY, CONF_PICTURE
 
 
-class TriggerEntity(update_coordinator.CoordinatorEntity):
+class TriggerEntity(CoordinatorEntity[TriggerUpdateCoordinator]):
     """Template entity based on trigger data."""
 
     domain: str

homeassistant/components/tibber/sensor.py

@@ -32,11 +32,14 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import PlatformNotReady
-from homeassistant.helpers import update_coordinator
 from homeassistant.helpers.device_registry import async_get as async_get_dev_reg
 from homeassistant.helpers.entity import DeviceInfo, EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import async_get as async_get_entity_reg
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
 from homeassistant.util import Throttle, dt as dt_util
 
 from .const import DOMAIN as TIBBER_DOMAIN, MANUFACTURER
@@ -239,7 +242,7 @@ async def async_setup_entry(
     entity_registry = async_get_entity_reg(hass)
     device_registry = async_get_dev_reg(hass)
 
-    coordinator: update_coordinator.DataUpdateCoordinator | None = None
+    coordinator: TibberDataCoordinator | None = None
     entities: list[TibberSensor] = []
     for home in tibber_connection.get_homes(only_active=False):
         try:
@@ -392,13 +395,13 @@ class TibberSensorElPrice(TibberSensor):
         ]["estimatedAnnualConsumption"]
 
 
-class TibberDataSensor(TibberSensor, update_coordinator.CoordinatorEntity):
+class TibberDataSensor(TibberSensor, CoordinatorEntity["TibberDataCoordinator"]):
     """Representation of a Tibber sensor."""
 
     def __init__(
         self,
         tibber_home,
-        coordinator: update_coordinator.DataUpdateCoordinator,
+        coordinator: TibberDataCoordinator,
         entity_description: SensorEntityDescription,
     ):
         """Initialize the sensor."""
@@ -420,7 +423,7 @@ class TibberDataSensor(TibberSensor, update_coordinator.CoordinatorEntity):
         return getattr(self._tibber_home, self.entity_description.key)
 
 
-class TibberSensorRT(TibberSensor, update_coordinator.CoordinatorEntity):
+class TibberSensorRT(TibberSensor, CoordinatorEntity["TibberRtDataCoordinator"]):
     """Representation of a Tibber sensor for real time consumption."""
 
     def __init__(
@@ -450,7 +453,7 @@ class TibberSensorRT(TibberSensor, update_coordinator.CoordinatorEntity):
 
     @callback
     def _handle_coordinator_update(self) -> None:
-        if not (live_measurement := self.coordinator.get_live_measurement()):  # type: ignore[attr-defined]
+        if not (live_measurement := self.coordinator.get_live_measurement()):
             return
         state = live_measurement.get(self.entity_description.key)
         if state is None:
@@ -479,7 +482,7 @@ class TibberSensorRT(TibberSensor, update_coordinator.CoordinatorEntity):
         self.async_write_ha_state()
 
 
-class TibberRtDataCoordinator(update_coordinator.DataUpdateCoordinator):
+class TibberRtDataCoordinator(DataUpdateCoordinator):
     """Handle Tibber realtime data."""
 
     def __init__(self, async_add_entities, tibber_home, hass):
@@ -538,7 +541,7 @@ class TibberRtDataCoordinator(update_coordinator.DataUpdateCoordinator):
         return self.data.get("data", {}).get("liveMeasurement")
 
 
-class TibberDataCoordinator(update_coordinator.DataUpdateCoordinator):
+class TibberDataCoordinator(DataUpdateCoordinator):
     """Handle Tibber data and insert statistics."""
 
     def __init__(self, hass, tibber_connection):

homeassistant/components/twentemilieu/entity.py

@@ -17,12 +17,14 @@ from homeassistant.helpers.update_coordinator import (
 from .const import DOMAIN
 
 
-class TwenteMilieuEntity(CoordinatorEntity[dict[WasteType, list[date]]], Entity):
+class TwenteMilieuEntity(
+    CoordinatorEntity[DataUpdateCoordinator[dict[WasteType, list[date]]]], Entity
+):
     """Defines a Twente Milieu entity."""
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator,
+        coordinator: DataUpdateCoordinator[dict[WasteType, list[date]]],
         entry: ConfigEntry,
     ) -> None:
         """Initialize the Twente Milieu entity."""

homeassistant/helpers/update_coordinator.py

@@ -6,7 +6,7 @@ from collections.abc import Awaitable, Callable
 from datetime import datetime, timedelta
 import logging
 from time import monotonic
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar  # pylint: disable=unused-import
 import urllib.error
 
 import aiohttp
@@ -24,6 +24,9 @@ REQUEST_REFRESH_DEFAULT_COOLDOWN = 10
 REQUEST_REFRESH_DEFAULT_IMMEDIATE = True
 
 _T = TypeVar("_T")
+_DataUpdateCoordinatorT = TypeVar(
+    "_DataUpdateCoordinatorT", bound="DataUpdateCoordinator[Any]"
+)
 
 
 class UpdateFailed(Exception):
@@ -295,10 +298,10 @@ class DataUpdateCoordinator(Generic[_T]):
             self._unsub_refresh = None
 
 
-class CoordinatorEntity(Generic[_T], entity.Entity):
+class CoordinatorEntity(entity.Entity, Generic[_DataUpdateCoordinatorT]):
     """A class for entities using DataUpdateCoordinator."""
 
-    def __init__(self, coordinator: DataUpdateCoordinator[_T]) -> None:
+    def __init__(self, coordinator: _DataUpdateCoordinatorT) -> None:
         """Create the entity with a DataUpdateCoordinator."""
         self.coordinator = coordinator