From 06db5476e4e687f1034196eb04c5cad3cf6861ad Mon Sep 17 00:00:00 2001 From: Marc Mueller <30130371+cdce8p@users.noreply.github.com> Date: Tue, 27 Dec 2022 22:47:04 +0100 Subject: [PATCH] Improve DataUpdateCoordinator typing in integrations (2) (#84656) --- homeassistant/components/atag/__init__.py | 8 +++++--- homeassistant/components/awair/__init__.py | 6 +++--- homeassistant/components/awair/sensor.py | 3 +-- homeassistant/components/brother/__init__.py | 2 +- homeassistant/components/canary/coordinator.py | 6 +++--- homeassistant/components/canary/model.py | 4 ++-- homeassistant/components/gree/climate.py | 5 +++-- homeassistant/components/mill/climate.py | 11 +++++++---- homeassistant/components/nam/__init__.py | 2 +- 9 files changed, 26 insertions(+), 21 deletions(-) diff --git a/homeassistant/components/atag/__init__.py b/homeassistant/components/atag/__init__.py index 56b2c1e969e..e8deca6f04d 100644 --- a/homeassistant/components/atag/__init__.py +++ b/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) diff --git a/homeassistant/components/awair/__init__.py b/homeassistant/components/awair/__init__.py index fd964328c4d..cef2c7d1fd4 100644 --- a/homeassistant/components/awair/__init__.py +++ b/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: diff --git a/homeassistant/components/awair/sensor.py b/homeassistant/components/awair/sensor.py index 18805154283..dc48e0f92c3 100644 --- a/homeassistant/components/awair/sensor.py +++ b/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: + if result := self.coordinator.data.get(self._device.uuid): return result.air_data return None diff --git a/homeassistant/components/brother/__init__.py b/homeassistant/components/brother/__init__.py index c1dbfd5bf0a..e0d0a0ca44c 100644 --- a/homeassistant/components/brother/__init__.py +++ b/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: diff --git a/homeassistant/components/canary/coordinator.py b/homeassistant/components/canary/coordinator.py index b2a8ef4daaa..d81589020e3 100644 --- a/homeassistant/components/canary/coordinator.py +++ b/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 diff --git a/homeassistant/components/canary/model.py b/homeassistant/components/canary/model.py index acf8f9c59a8..4ed868a7e60 100644 --- a/homeassistant/components/canary/model.py +++ b/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]] diff --git a/homeassistant/components/gree/climate.py b/homeassistant/components/gree/climate.py index 6afbca7789b..8a53e3b3229 100644 --- a/homeassistant/components/gree/climate.py +++ b/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 diff --git a/homeassistant/components/mill/climate.py b/homeassistant/components/mill/climate.py index 912495d05e8..42b759b3cdf 100644 --- a/homeassistant/components/mill/climate.py +++ b/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 diff --git a/homeassistant/components/nam/__init__.py b/homeassistant/components/nam/__init__.py index 0fbc9384634..c035184c37b 100644 --- a/homeassistant/components/nam/__init__.py +++ b/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__(