Improve DataUpdateCoordinator typing in integrations (5) (#84740)

pull/84747/head
Marc Mueller 2022-12-28 23:44:41 +01:00 committed by GitHub
parent 5c43f0861f
commit 2512eb1e4c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 40 additions and 17 deletions

View File

@ -63,7 +63,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Add cert-expiry entry."""
coordinator = hass.data[DOMAIN][entry.entry_id]
coordinator: CertExpiryDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
sensors = [
SSLCertificateTimestamp(coordinator),
@ -72,7 +72,7 @@ async def async_setup_entry(
async_add_entities(sensors, True)
class CertExpiryEntity(CoordinatorEntity):
class CertExpiryEntity(CoordinatorEntity[CertExpiryDataUpdateCoordinator]):
"""Defines a base Cert Expiry entity."""
_attr_icon = "mdi:certificate"

View File

@ -1,6 +1,7 @@
"""The Huisbaasje integration."""
from datetime import timedelta
import logging
from typing import Any
import async_timeout
from energyflip import EnergyFlip, EnergyFlipException
@ -45,7 +46,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
_LOGGER.error("Authentication failed: %s", str(exception))
return False
async def async_update_data():
async def async_update_data() -> dict[str, dict[str, Any]]:
return await async_update_huisbaasje(energyflip)
# Create a coordinator for polling updates
@ -80,7 +81,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return unload_ok
async def async_update_huisbaasje(energyflip):
async def async_update_huisbaasje(energyflip: EnergyFlip) -> dict[str, dict[str, Any]]:
"""Update the data by performing a request to Huisbaasje."""
try:
# Note: asyncio.TimeoutError and aiohttp.ClientError are already

View File

@ -3,6 +3,7 @@ from __future__ import annotations
from dataclasses import dataclass
import logging
from typing import Any
from energyflip.const import (
SOURCE_TYPE_ELECTRICITY,
@ -234,7 +235,9 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up the sensor platform."""
coordinator = hass.data[DOMAIN][config_entry.entry_id][DATA_COORDINATOR]
coordinator: DataUpdateCoordinator[dict[str, dict[str, Any]]] = hass.data[DOMAIN][
config_entry.entry_id
][DATA_COORDINATOR]
user_id = config_entry.data[CONF_ID]
async_add_entities(
@ -243,14 +246,16 @@ async def async_setup_entry(
)
class HuisbaasjeSensor(CoordinatorEntity, SensorEntity):
class HuisbaasjeSensor(
CoordinatorEntity[DataUpdateCoordinator[dict[str, dict[str, Any]]]], SensorEntity
):
"""Defines a Huisbaasje sensor."""
entity_description: HuisbaasjeSensorEntityDescription
def __init__(
self,
coordinator: DataUpdateCoordinator,
coordinator: DataUpdateCoordinator[dict[str, dict[str, Any]]],
user_id: str,
description: HuisbaasjeSensorEntityDescription,
) -> None:
@ -278,7 +283,7 @@ class HuisbaasjeSensor(CoordinatorEntity, SensorEntity):
@property
def available(self) -> bool:
"""Return if entity is available."""
return (
return bool(
super().available
and self.coordinator.data
and self._source_type in self.coordinator.data

View File

@ -5,6 +5,7 @@ from datetime import timedelta
import logging
import ultraheat_api
from ultraheat_api.response import HeatMeterResponse
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_DEVICE, Platform
@ -26,7 +27,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
reader = ultraheat_api.UltraheatReader(entry.data[CONF_DEVICE])
api = ultraheat_api.HeatMeterService(reader)
async def async_update_data():
async def async_update_data() -> HeatMeterResponse:
"""Fetch data from the API."""
_LOGGER.debug("Polling on %s", entry.data[CONF_DEVICE])
return await hass.async_add_executor_job(api.read)

View File

@ -4,12 +4,21 @@ from __future__ import annotations
from dataclasses import asdict
import logging
from homeassistant.components.sensor import RestoreSensor, SensorDeviceClass
from ultraheat_api.response import HeatMeterResponse
from homeassistant.components.sensor import (
RestoreSensor,
SensorDeviceClass,
SensorEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
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 homeassistant.util import dt as dt_util
from . import DOMAIN
@ -23,7 +32,9 @@ async def async_setup_entry(
) -> None:
"""Set up the sensor platform."""
unique_id = entry.entry_id
coordinator = hass.data[DOMAIN][entry.entry_id]
coordinator: DataUpdateCoordinator[HeatMeterResponse] = hass.data[DOMAIN][
entry.entry_id
]
model = entry.data["model"]
@ -42,16 +53,21 @@ async def async_setup_entry(
async_add_entities(sensors)
class HeatMeterSensor(CoordinatorEntity, RestoreSensor):
class HeatMeterSensor(
CoordinatorEntity[DataUpdateCoordinator[HeatMeterResponse]], RestoreSensor
):
"""Representation of a Sensor."""
def __init__(self, coordinator, description, device):
def __init__(
self,
coordinator: DataUpdateCoordinator[HeatMeterResponse],
description: SensorEntityDescription,
device: DeviceInfo,
) -> None:
"""Set up the sensor with the initial values."""
super().__init__(coordinator)
self.key = description.key
self._attr_unique_id = (
f"{coordinator.config_entry.data['device_number']}_{description.key}"
)
self._attr_unique_id = f"{coordinator.config_entry.data['device_number']}_{description.key}" # type: ignore[union-attr]
self._attr_name = f"Heat Meter {description.name}"
self.entity_description = description