Improve iqvia typing (#84734)

This commit is contained in:
Marc Mueller 2022-12-29 09:45:05 +01:00 committed by GitHub
parent c4b45fb110
commit bfb509ccb8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 13 additions and 12 deletions

View File

@ -2,10 +2,10 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
from collections.abc import Awaitable, Callable from collections.abc import Callable, Coroutine
from datetime import timedelta from datetime import timedelta
from functools import partial from functools import partial
from typing import Any, cast from typing import Any
from pyiqvia import Client from pyiqvia import Client
from pyiqvia.errors import IQVIAError from pyiqvia.errors import IQVIAError
@ -57,16 +57,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
client.disable_request_retries() client.disable_request_retries()
async def async_get_data_from_api( async def async_get_data_from_api(
api_coro: Callable[..., Awaitable] api_coro: Callable[..., Coroutine[Any, Any, dict[str, Any]]]
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Get data from a particular API coroutine.""" """Get data from a particular API coroutine."""
try: try:
data = await api_coro() return await api_coro()
except IQVIAError as err: except IQVIAError as err:
raise UpdateFailed from err raise UpdateFailed from err
return cast(dict[str, Any], data)
coordinators = {} coordinators = {}
init_data_update_tasks = [] init_data_update_tasks = []
@ -115,14 +113,14 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return unload_ok return unload_ok
class IQVIAEntity(CoordinatorEntity): class IQVIAEntity(CoordinatorEntity[DataUpdateCoordinator[dict[str, Any]]]):
"""Define a base IQVIA entity.""" """Define a base IQVIA entity."""
_attr_has_entity_name = True _attr_has_entity_name = True
def __init__( def __init__(
self, self,
coordinator: DataUpdateCoordinator, coordinator: DataUpdateCoordinator[dict[str, Any]],
entry: ConfigEntry, entry: ConfigEntry,
description: EntityDescription, description: EntityDescription,
) -> None: ) -> None:

View File

@ -35,7 +35,9 @@ async def async_get_config_entry_diagnostics(
hass: HomeAssistant, entry: ConfigEntry hass: HomeAssistant, entry: ConfigEntry
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Return diagnostics for a config entry.""" """Return diagnostics for a config entry."""
coordinators: dict[str, DataUpdateCoordinator] = hass.data[DOMAIN][entry.entry_id] coordinators: dict[str, DataUpdateCoordinator[dict[str, Any]]] = hass.data[DOMAIN][
entry.entry_id
]
return { return {
"entry": async_redact_data(entry.as_dict(), TO_REDACT), "entry": async_redact_data(entry.as_dict(), TO_REDACT),

View File

@ -2,7 +2,7 @@
from __future__ import annotations from __future__ import annotations
from statistics import mean from statistics import mean
from typing import NamedTuple from typing import Any, NamedTuple, cast
import numpy as np import numpy as np
@ -247,10 +247,11 @@ class IndexSensor(IQVIAEntity, SensorEntity):
key = self.entity_description.key.split("_")[-1].title() key = self.entity_description.key.split("_")[-1].title()
try: try:
[period] = [p for p in data["periods"] if p["Type"] == key] [period] = [p for p in data["periods"] if p["Type"] == key] # type: ignore[index]
except ValueError: except TypeError:
return return
data = cast(dict[str, Any], data)
[rating] = [ [rating] = [
i.label for i in RATING_MAPPING if i.minimum <= period["Index"] <= i.maximum i.label for i in RATING_MAPPING if i.minimum <= period["Index"] <= i.maximum
] ]