Use EntityDescription - iqvia (#55218)

This commit is contained in:
Aaron Bach 2021-09-12 09:34:44 -06:00 committed by GitHub
parent f1556ead6d
commit 4e8db7173a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 101 additions and 50 deletions

View File

@ -10,12 +10,12 @@ from typing import Any, Callable, Dict, cast
from pyiqvia import Client
from pyiqvia.errors import IQVIAError
from homeassistant.components.sensor import SensorEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_ATTRIBUTION
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import ConfigEntryNotReady
from homeassistant.helpers import aiohttp_client
from homeassistant.helpers.entity import EntityDescription
from homeassistant.helpers.update_coordinator import (
CoordinatorEntity,
DataUpdateCoordinator,
@ -107,27 +107,22 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return unload_ok
class IQVIAEntity(CoordinatorEntity, SensorEntity):
class IQVIAEntity(CoordinatorEntity):
"""Define a base IQVIA entity."""
def __init__(
self,
coordinator: DataUpdateCoordinator,
entry: ConfigEntry,
sensor_type: str,
name: str,
icon: str,
description: EntityDescription,
) -> None:
"""Initialize."""
super().__init__(coordinator)
self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: DEFAULT_ATTRIBUTION}
self._attr_icon = icon
self._attr_name = name
self._attr_unique_id = f"{entry.data[CONF_ZIP_CODE]}_{sensor_type}"
self._attr_native_unit_of_measurement = "index"
self._attr_unique_id = f"{entry.data[CONF_ZIP_CODE]}_{description.key}"
self._entry = entry
self._type = sensor_type
self.entity_description = description
@callback
def _handle_coordinator_update(self) -> None:
@ -142,7 +137,7 @@ class IQVIAEntity(CoordinatorEntity, SensorEntity):
"""Register callbacks."""
await super().async_added_to_hass()
if self._type == TYPE_ALLERGY_FORECAST:
if self.entity_description.key == TYPE_ALLERGY_FORECAST:
self.async_on_remove(
self.hass.data[DOMAIN][DATA_COORDINATOR][self._entry.entry_id][
TYPE_ALLERGY_OUTLOOK

View File

@ -21,14 +21,3 @@ TYPE_ASTHMA_TOMORROW = "asthma_index_tomorrow"
TYPE_DISEASE_FORECAST = "disease_average_forecasted"
TYPE_DISEASE_INDEX = "disease_index"
TYPE_DISEASE_TODAY = "disease_index_today"
SENSORS = {
TYPE_ALLERGY_FORECAST: ("Allergy Index: Forecasted Average", "mdi:flower"),
TYPE_ALLERGY_TODAY: ("Allergy Index: Today", "mdi:flower"),
TYPE_ALLERGY_TOMORROW: ("Allergy Index: Tomorrow", "mdi:flower"),
TYPE_ASTHMA_FORECAST: ("Asthma Index: Forecasted Average", "mdi:flower"),
TYPE_ASTHMA_TODAY: ("Asthma Index: Today", "mdi:flower"),
TYPE_ASTHMA_TOMORROW: ("Asthma Index: Tomorrow", "mdi:flower"),
TYPE_DISEASE_FORECAST: ("Cold & Flu: Forecasted Average", "mdi:snowflake"),
TYPE_DISEASE_TODAY: ("Cold & Flu Index: Today", "mdi:pill"),
}

View File

@ -5,6 +5,11 @@ from statistics import mean
import numpy as np
from homeassistant.components.sensor import (
STATE_CLASS_MEASUREMENT,
SensorEntity,
SensorEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_STATE
from homeassistant.core import HomeAssistant, callback
@ -14,7 +19,6 @@ from . import IQVIAEntity
from .const import (
DATA_COORDINATOR,
DOMAIN,
SENSORS,
TYPE_ALLERGY_FORECAST,
TYPE_ALLERGY_INDEX,
TYPE_ALLERGY_OUTLOOK,
@ -57,32 +61,89 @@ RATING_MAPPING = [
{"label": "High", "minimum": 9.7, "maximum": 12},
]
TREND_FLAT = "Flat"
TREND_INCREASING = "Increasing"
TREND_SUBSIDING = "Subsiding"
FORECAST_SENSOR_DESCRIPTIONS = (
SensorEntityDescription(
key=TYPE_ALLERGY_FORECAST,
name="Allergy Index: Forecasted Average",
icon="mdi:flower",
),
SensorEntityDescription(
key=TYPE_ASTHMA_FORECAST,
name="Asthma Index: Forecasted Average",
icon="mdi:flower",
),
SensorEntityDescription(
key=TYPE_DISEASE_FORECAST,
name="Cold & Flu: Forecasted Average",
icon="mdi:snowflake",
),
)
INDEX_SENSOR_DESCRIPTIONS = (
SensorEntityDescription(
key=TYPE_ALLERGY_TODAY,
name="Allergy Index: Today",
icon="mdi:flower",
state_class=STATE_CLASS_MEASUREMENT,
),
SensorEntityDescription(
key=TYPE_ALLERGY_TOMORROW,
name="Allergy Index: Tomorrow",
icon="mdi:flower",
),
SensorEntityDescription(
key=TYPE_ASTHMA_TODAY,
name="Asthma Index: Today",
icon="mdi:flower",
state_class=STATE_CLASS_MEASUREMENT,
),
SensorEntityDescription(
key=TYPE_ASTHMA_TOMORROW,
name="Asthma Index: Tomorrow",
icon="mdi:flower",
state_class=STATE_CLASS_MEASUREMENT,
),
SensorEntityDescription(
key=TYPE_DISEASE_TODAY,
name="Cold & Flu Index: Today",
icon="mdi:pill",
state_class=STATE_CLASS_MEASUREMENT,
),
)
async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None:
"""Set up IQVIA sensors based on a config entry."""
sensor_class_mapping = {
TYPE_ALLERGY_FORECAST: ForecastSensor,
TYPE_ALLERGY_TODAY: IndexSensor,
TYPE_ALLERGY_TOMORROW: IndexSensor,
TYPE_ASTHMA_FORECAST: ForecastSensor,
TYPE_ASTHMA_TODAY: IndexSensor,
TYPE_ASTHMA_TOMORROW: IndexSensor,
TYPE_DISEASE_FORECAST: ForecastSensor,
TYPE_DISEASE_TODAY: IndexSensor,
}
sensors = []
for sensor_type, (name, icon) in SENSORS.items():
api_category = API_CATEGORY_MAPPING.get(sensor_type, sensor_type)
coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][api_category]
sensor_class = sensor_class_mapping[sensor_type]
sensors.append(sensor_class(coordinator, entry, sensor_type, name, icon))
sensors: list[ForecastSensor | IndexSensor] = [
ForecastSensor(
hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][
API_CATEGORY_MAPPING.get(description.key, description.key)
],
entry,
description,
)
for description in FORECAST_SENSOR_DESCRIPTIONS
]
sensors.extend(
[
IndexSensor(
hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][
API_CATEGORY_MAPPING.get(description.key, description.key)
],
entry,
description,
)
for description in INDEX_SENSOR_DESCRIPTIONS
]
)
async_add_entities(sensors)
@ -104,7 +165,7 @@ def calculate_trend(indices: list[float]) -> str:
return TREND_FLAT
class ForecastSensor(IQVIAEntity):
class ForecastSensor(IQVIAEntity, SensorEntity):
"""Define sensor related to forecast data."""
@callback
@ -137,7 +198,7 @@ class ForecastSensor(IQVIAEntity):
}
)
if self._type == TYPE_ALLERGY_FORECAST:
if self.entity_description.key == TYPE_ALLERGY_FORECAST:
outlook_coordinator = self.hass.data[DOMAIN][DATA_COORDINATOR][
self._entry.entry_id
][TYPE_ALLERGY_OUTLOOK]
@ -153,7 +214,7 @@ class ForecastSensor(IQVIAEntity):
] = outlook_coordinator.data.get("Season")
class IndexSensor(IQVIAEntity):
class IndexSensor(IQVIAEntity, SensorEntity):
"""Define sensor related to indices."""
@callback
@ -163,16 +224,22 @@ class IndexSensor(IQVIAEntity):
return
try:
if self._type in (TYPE_ALLERGY_TODAY, TYPE_ALLERGY_TOMORROW):
if self.entity_description.key in (
TYPE_ALLERGY_TODAY,
TYPE_ALLERGY_TOMORROW,
):
data = self.coordinator.data.get("Location")
elif self._type in (TYPE_ASTHMA_TODAY, TYPE_ASTHMA_TOMORROW):
elif self.entity_description.key in (
TYPE_ASTHMA_TODAY,
TYPE_ASTHMA_TOMORROW,
):
data = self.coordinator.data.get("Location")
elif self._type == TYPE_DISEASE_TODAY:
elif self.entity_description.key == TYPE_DISEASE_TODAY:
data = self.coordinator.data.get("Location")
except KeyError:
return
key = self._type.split("_")[-1].title()
key = self.entity_description.key.split("_")[-1].title()
try:
[period] = [p for p in data["periods"] if p["Type"] == key]
@ -194,7 +261,7 @@ class IndexSensor(IQVIAEntity):
}
)
if self._type in (TYPE_ALLERGY_TODAY, TYPE_ALLERGY_TOMORROW):
if self.entity_description.key in (TYPE_ALLERGY_TODAY, TYPE_ALLERGY_TOMORROW):
for idx, attrs in enumerate(period["Triggers"]):
index = idx + 1
self._attr_extra_state_attributes.update(
@ -204,7 +271,7 @@ class IndexSensor(IQVIAEntity):
f"{ATTR_ALLERGEN_TYPE}_{index}": attrs["PlantType"],
}
)
elif self._type in (TYPE_ASTHMA_TODAY, TYPE_ASTHMA_TOMORROW):
elif self.entity_description.key in (TYPE_ASTHMA_TODAY, TYPE_ASTHMA_TOMORROW):
for idx, attrs in enumerate(period["Triggers"]):
index = idx + 1
self._attr_extra_state_attributes.update(
@ -213,7 +280,7 @@ class IndexSensor(IQVIAEntity):
f"{ATTR_ALLERGEN_AMOUNT}_{index}": attrs["PPM"],
}
)
elif self._type == TYPE_DISEASE_TODAY:
elif self.entity_description.key == TYPE_DISEASE_TODAY:
for attrs in period["Triggers"]:
self._attr_extra_state_attributes[
f"{attrs['Name'].lower()}_index"