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

This commit is contained in:
Marc Mueller 2022-12-27 22:47:04 +01:00 committed by GitHub
parent 1de41ab123
commit 06db5476e4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 26 additions and 21 deletions

View file

@ -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)

View file

@ -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:

View file

@ -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

View file

@ -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:

View file

@ -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

View file

@ -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]]

View file

@ -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

View file

@ -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

View file

@ -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__(