Improve discovergy generic typing (#94131)

pull/94136/head
Marc Mueller 2023-06-07 01:47:13 +02:00 committed by GitHub
parent deaf910a24
commit 4c80420d2d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 12 additions and 11 deletions

View File

@ -7,7 +7,7 @@ import logging
import pydiscovergy
from pydiscovergy.authentication import BasicAuth
import pydiscovergy.error as discovergyError
from pydiscovergy.models import Meter
from pydiscovergy.models import Meter, Reading
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, Platform
@ -29,7 +29,9 @@ class DiscovergyData:
api_client: pydiscovergy.Discovergy = field(default_factory=lambda: None)
meters: list[Meter] = field(default_factory=lambda: [])
coordinators: dict[str, DataUpdateCoordinator] = field(default_factory=lambda: {})
coordinators: dict[str, DataUpdateCoordinator[Reading]] = field(
default_factory=lambda: {}
)
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:

View File

@ -9,7 +9,6 @@ from homeassistant.components.diagnostics import async_redact_data
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, CONF_UNIQUE_ID
from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from . import DiscovergyData
from .const import DOMAIN
@ -40,7 +39,7 @@ async def async_get_config_entry_diagnostics(
flattened_meter.append(async_redact_data(meter.__dict__, TO_REDACT_METER))
# get last reading for meter and make a dict of it
coordinator: DataUpdateCoordinator = data.coordinators[meter.get_meter_id()]
coordinator = data.coordinators[meter.get_meter_id()]
last_readings[meter.get_meter_id()] = coordinator.data.__dict__
return {

View File

@ -5,7 +5,7 @@ import logging
from pydiscovergy import Discovergy
from pydiscovergy.error import AccessTokenExpired, HTTPError
from pydiscovergy.models import Meter
from pydiscovergy.models import Meter, Reading
from homeassistant.components.sensor import (
SensorDeviceClass,
@ -50,7 +50,7 @@ _LOGGER = logging.getLogger(__name__)
class DiscovergyMixin:
"""Mixin for alternative keys."""
alternative_keys: list = field(default_factory=lambda: [])
alternative_keys: list[str] = field(default_factory=lambda: [])
scale: int = field(default_factory=lambda: 1000)
@ -165,10 +165,10 @@ def get_coordinator_for_meter(
meter: Meter,
discovergy_instance: Discovergy,
update_interval: timedelta,
) -> DataUpdateCoordinator:
) -> DataUpdateCoordinator[Reading]:
"""Create a new DataUpdateCoordinator for given meter."""
async def async_update_data():
async def async_update_data() -> Reading:
"""Fetch data from API endpoint."""
try:
return await discovergy_instance.get_last_reading(meter.get_meter_id())
@ -205,7 +205,7 @@ async def async_setup_entry(
seconds=entry.options.get(CONF_TIME_BETWEEN_UPDATE, DEFAULT_TIME_BETWEEN_UPDATE)
)
entities = []
entities: list[DiscovergySensor] = []
for meter in meters:
# Get coordinator for meter, set config entry and fetch initial data
# so we have data when entities are added
@ -238,7 +238,7 @@ async def async_setup_entry(
async_add_entities(entities, False)
class DiscovergySensor(CoordinatorEntity, SensorEntity):
class DiscovergySensor(CoordinatorEntity[DataUpdateCoordinator[Reading]], SensorEntity):
"""Represents a discovergy smart meter sensor."""
entity_description: DiscovergySensorEntityDescription
@ -250,7 +250,7 @@ class DiscovergySensor(CoordinatorEntity, SensorEntity):
data_key: str,
description: DiscovergySensorEntityDescription,
meter: Meter,
coordinator: DataUpdateCoordinator,
coordinator: DataUpdateCoordinator[Reading],
) -> None:
"""Initialize the sensor."""
super().__init__(coordinator)