Improve DataUpdateCoordinator typing in integrations (2) (#84656)

homeassistant/components/atag/__init__.py

@@ -37,7 +37,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     atag = AtagOne(
         session=async_get_clientsession(hass), **entry.data, device=entry.unique_id
     )
-    coordinator = DataUpdateCoordinator(
+    coordinator = DataUpdateCoordinator[AtagOne](
         hass,
         _LOGGER,
         name=DOMAIN.title(),
@@ -65,10 +65,12 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return unload_ok
 
 
-class AtagEntity(CoordinatorEntity):
+class AtagEntity(CoordinatorEntity[DataUpdateCoordinator[AtagOne]]):
     """Defines a base Atag entity."""
 
-    def __init__(self, coordinator: DataUpdateCoordinator, atag_id: str) -> None:
+    def __init__(
+        self, coordinator: DataUpdateCoordinator[AtagOne], atag_id: str
+    ) -> None:
         """Initialize the Atag entity."""
         super().__init__(coordinator)
 

homeassistant/components/awair/__init__.py

@@ -72,7 +72,7 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
     return unload_ok
 
 
-class AwairDataUpdateCoordinator(DataUpdateCoordinator):
+class AwairDataUpdateCoordinator(DataUpdateCoordinator[dict[str, AwairResult]]):
     """Define a wrapper class to update Awair data."""
 
     def __init__(
@@ -107,7 +107,7 @@ class AwairCloudDataUpdateCoordinator(AwairDataUpdateCoordinator):
 
         super().__init__(hass, config_entry, UPDATE_INTERVAL_CLOUD)
 
-    async def _async_update_data(self) -> dict[str, AwairResult] | None:
+    async def _async_update_data(self) -> dict[str, AwairResult]:
         """Update data via Awair client library."""
         async with timeout(API_TIMEOUT):
             try:
@@ -139,7 +139,7 @@ class AwairLocalDataUpdateCoordinator(AwairDataUpdateCoordinator):
 
         super().__init__(hass, config_entry, UPDATE_INTERVAL_LOCAL)
 
-    async def _async_update_data(self) -> dict[str, AwairResult] | None:
+    async def _async_update_data(self) -> dict[str, AwairResult]:
         """Update data via Awair client library."""
         async with timeout(API_TIMEOUT):
             try:

homeassistant/components/awair/sensor.py

@@ -215,8 +215,7 @@ class AwairSensor(CoordinatorEntity[AwairDataUpdateCoordinator], SensorEntity):
     @property
     def _air_data(self) -> AirData | None:
         """Return the latest data for our device, or None."""
-        result: AwairResult | None = self.coordinator.data.get(self._device.uuid)
+        if result := self.coordinator.data.get(self._device.uuid):
-        if result:
             return result.air_data
 
         return None

homeassistant/components/brother/__init__.py

@@ -62,7 +62,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return unload_ok
 
 
-class BrotherDataUpdateCoordinator(DataUpdateCoordinator):
+class BrotherDataUpdateCoordinator(DataUpdateCoordinator[BrotherSensors]):
     """Class to manage fetching Brother data from the printer."""
 
     def __init__(self, hass: HomeAssistant, brother: Brother) -> None:

homeassistant/components/canary/coordinator.py

@@ -7,7 +7,7 @@ import logging
 
 from async_timeout import timeout
 from canary.api import Api
-from canary.model import Location
+from canary.model import Location, Reading
 from requests.exceptions import ConnectTimeout, HTTPError
 
 from homeassistant.core import HomeAssistant
@@ -19,7 +19,7 @@ from .model import CanaryData
 _LOGGER = logging.getLogger(__name__)
 
 
-class CanaryDataUpdateCoordinator(DataUpdateCoordinator):
+class CanaryDataUpdateCoordinator(DataUpdateCoordinator[CanaryData]):
     """Class to manage fetching Canary data."""
 
     def __init__(self, hass: HomeAssistant, *, api: Api) -> None:
@@ -37,7 +37,7 @@ class CanaryDataUpdateCoordinator(DataUpdateCoordinator):
     def _update_data(self) -> CanaryData:
         """Fetch data from Canary via sync functions."""
         locations_by_id: dict[str, Location] = {}
-        readings_by_device_id: dict[str, ValuesView] = {}
+        readings_by_device_id: dict[str, ValuesView[Reading]] = {}
 
         for location in self.canary.get_locations():
             location_id = location.location_id

homeassistant/components/canary/model.py

@@ -4,11 +4,11 @@ from __future__ import annotations
 from collections.abc import ValuesView
 from typing import TypedDict
 
-from canary.model import Location
+from canary.model import Location, Reading
 
 
 class CanaryData(TypedDict):
     """TypedDict for Canary Coordinator Data."""
 
     locations: dict[str, Location]
-    readings: dict[str, ValuesView]
+    readings: dict[str, ValuesView[Reading]]

homeassistant/components/gree/climate.py

@@ -43,6 +43,7 @@ from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from .bridge import DeviceDataUpdateCoordinator
 from .const import (
     COORDINATORS,
     DISPATCH_DEVICE_DISCOVERED,
@@ -105,7 +106,7 @@ async def async_setup_entry(
     )
 
 
-class GreeClimateEntity(CoordinatorEntity, ClimateEntity):
+class GreeClimateEntity(CoordinatorEntity[DeviceDataUpdateCoordinator], ClimateEntity):
     """Representation of a Gree HVAC device."""
 
     _attr_precision = PRECISION_WHOLE
@@ -116,7 +117,7 @@ class GreeClimateEntity(CoordinatorEntity, ClimateEntity):
         | ClimateEntityFeature.SWING_MODE
     )
 
-    def __init__(self, coordinator):
+    def __init__(self, coordinator: DeviceDataUpdateCoordinator) -> None:
         """Initialize the Gree device."""
         super().__init__(coordinator)
         self._name = coordinator.device.device_info.name

homeassistant/components/mill/climate.py

@@ -28,6 +28,7 @@ from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import MillDataUpdateCoordinator
 from .const import (
     ATTR_AWAY_TEMP,
     ATTR_COMFORT_TEMP,
@@ -86,7 +87,7 @@ async def async_setup_entry(
     )
 
 
-class MillHeater(CoordinatorEntity, ClimateEntity):
+class MillHeater(CoordinatorEntity[MillDataUpdateCoordinator], ClimateEntity):
     """Representation of a Mill Thermostat device."""
 
     _attr_fan_modes = [FAN_ON, FAN_OFF]
@@ -94,7 +95,9 @@ class MillHeater(CoordinatorEntity, ClimateEntity):
     _attr_min_temp = MIN_TEMP
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
 
-    def __init__(self, coordinator, heater):
+    def __init__(
+        self, coordinator: MillDataUpdateCoordinator, heater: mill.Heater
+    ) -> None:
         """Initialize the thermostat."""
 
         super().__init__(coordinator)
@@ -196,7 +199,7 @@ class MillHeater(CoordinatorEntity, ClimateEntity):
             self._attr_hvac_mode = HVACMode.OFF
 
 
-class LocalMillHeater(CoordinatorEntity, ClimateEntity):
+class LocalMillHeater(CoordinatorEntity[MillDataUpdateCoordinator], ClimateEntity):
     """Representation of a Mill Thermostat device."""
 
     _attr_hvac_mode = HVACMode.HEAT
@@ -207,7 +210,7 @@ class LocalMillHeater(CoordinatorEntity, ClimateEntity):
     _attr_target_temperature_step = PRECISION_HALVES
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
 
-    def __init__(self, coordinator):
+    def __init__(self, coordinator: MillDataUpdateCoordinator) -> None:
         """Initialize the thermostat."""
         super().__init__(coordinator)
         self._attr_name = coordinator.mill_data_connection.name

homeassistant/components/nam/__init__.py

@@ -92,7 +92,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return unload_ok
 
 
-class NAMDataUpdateCoordinator(DataUpdateCoordinator):
+class NAMDataUpdateCoordinator(DataUpdateCoordinator[NAMSensors]):
     """Class to manage fetching Nettigo Air Monitor data."""
 
     def __init__(