Improve discovergy generic typing (#94131)

This commit is contained in:
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 import pydiscovergy
from pydiscovergy.authentication import BasicAuth from pydiscovergy.authentication import BasicAuth
import pydiscovergy.error as discovergyError import pydiscovergy.error as discovergyError
from pydiscovergy.models import Meter from pydiscovergy.models import Meter, Reading
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, Platform from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, Platform
@ -29,7 +29,9 @@ class DiscovergyData:
api_client: pydiscovergy.Discovergy = field(default_factory=lambda: None) api_client: pydiscovergy.Discovergy = field(default_factory=lambda: None)
meters: list[Meter] = field(default_factory=lambda: []) 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: 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.config_entries import ConfigEntry
from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, CONF_UNIQUE_ID from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, CONF_UNIQUE_ID
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from . import DiscovergyData from . import DiscovergyData
from .const import DOMAIN 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)) flattened_meter.append(async_redact_data(meter.__dict__, TO_REDACT_METER))
# get last reading for meter and make a dict of it # 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__ last_readings[meter.get_meter_id()] = coordinator.data.__dict__
return { return {

View File

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