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

View File

@ -21,14 +21,3 @@ TYPE_ASTHMA_TOMORROW = "asthma_index_tomorrow"
TYPE_DISEASE_FORECAST = "disease_average_forecasted" TYPE_DISEASE_FORECAST = "disease_average_forecasted"
TYPE_DISEASE_INDEX = "disease_index" TYPE_DISEASE_INDEX = "disease_index"
TYPE_DISEASE_TODAY = "disease_index_today" 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 import numpy as np
from homeassistant.components.sensor import (
STATE_CLASS_MEASUREMENT,
SensorEntity,
SensorEntityDescription,
)
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_STATE from homeassistant.const import ATTR_STATE
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
@ -14,7 +19,6 @@ from . import IQVIAEntity
from .const import ( from .const import (
DATA_COORDINATOR, DATA_COORDINATOR,
DOMAIN, DOMAIN,
SENSORS,
TYPE_ALLERGY_FORECAST, TYPE_ALLERGY_FORECAST,
TYPE_ALLERGY_INDEX, TYPE_ALLERGY_INDEX,
TYPE_ALLERGY_OUTLOOK, TYPE_ALLERGY_OUTLOOK,
@ -57,32 +61,89 @@ RATING_MAPPING = [
{"label": "High", "minimum": 9.7, "maximum": 12}, {"label": "High", "minimum": 9.7, "maximum": 12},
] ]
TREND_FLAT = "Flat" TREND_FLAT = "Flat"
TREND_INCREASING = "Increasing" TREND_INCREASING = "Increasing"
TREND_SUBSIDING = "Subsiding" 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( async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None: ) -> None:
"""Set up IQVIA sensors based on a config entry.""" """Set up IQVIA sensors based on a config entry."""
sensor_class_mapping = { sensors: list[ForecastSensor | IndexSensor] = [
TYPE_ALLERGY_FORECAST: ForecastSensor, ForecastSensor(
TYPE_ALLERGY_TODAY: IndexSensor, hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][
TYPE_ALLERGY_TOMORROW: IndexSensor, API_CATEGORY_MAPPING.get(description.key, description.key)
TYPE_ASTHMA_FORECAST: ForecastSensor, ],
TYPE_ASTHMA_TODAY: IndexSensor, entry,
TYPE_ASTHMA_TOMORROW: IndexSensor, description,
TYPE_DISEASE_FORECAST: ForecastSensor, )
TYPE_DISEASE_TODAY: IndexSensor, for description in FORECAST_SENSOR_DESCRIPTIONS
} ]
sensors.extend(
sensors = [] [
for sensor_type, (name, icon) in SENSORS.items(): IndexSensor(
api_category = API_CATEGORY_MAPPING.get(sensor_type, sensor_type) hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][
coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][api_category] API_CATEGORY_MAPPING.get(description.key, description.key)
sensor_class = sensor_class_mapping[sensor_type] ],
sensors.append(sensor_class(coordinator, entry, sensor_type, name, icon)) entry,
description,
)
for description in INDEX_SENSOR_DESCRIPTIONS
]
)
async_add_entities(sensors) async_add_entities(sensors)
@ -104,7 +165,7 @@ def calculate_trend(indices: list[float]) -> str:
return TREND_FLAT return TREND_FLAT
class ForecastSensor(IQVIAEntity): class ForecastSensor(IQVIAEntity, SensorEntity):
"""Define sensor related to forecast data.""" """Define sensor related to forecast data."""
@callback @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][ outlook_coordinator = self.hass.data[DOMAIN][DATA_COORDINATOR][
self._entry.entry_id self._entry.entry_id
][TYPE_ALLERGY_OUTLOOK] ][TYPE_ALLERGY_OUTLOOK]
@ -153,7 +214,7 @@ class ForecastSensor(IQVIAEntity):
] = outlook_coordinator.data.get("Season") ] = outlook_coordinator.data.get("Season")
class IndexSensor(IQVIAEntity): class IndexSensor(IQVIAEntity, SensorEntity):
"""Define sensor related to indices.""" """Define sensor related to indices."""
@callback @callback
@ -163,16 +224,22 @@ class IndexSensor(IQVIAEntity):
return return
try: 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") 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") 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") data = self.coordinator.data.get("Location")
except KeyError: except KeyError:
return return
key = self._type.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]
@ -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"]): for idx, attrs in enumerate(period["Triggers"]):
index = idx + 1 index = idx + 1
self._attr_extra_state_attributes.update( self._attr_extra_state_attributes.update(
@ -204,7 +271,7 @@ class IndexSensor(IQVIAEntity):
f"{ATTR_ALLERGEN_TYPE}_{index}": attrs["PlantType"], 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"]): for idx, attrs in enumerate(period["Triggers"]):
index = idx + 1 index = idx + 1
self._attr_extra_state_attributes.update( self._attr_extra_state_attributes.update(
@ -213,7 +280,7 @@ class IndexSensor(IQVIAEntity):
f"{ATTR_ALLERGEN_AMOUNT}_{index}": attrs["PPM"], 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"]: for attrs in period["Triggers"]:
self._attr_extra_state_attributes[ self._attr_extra_state_attributes[
f"{attrs['Name'].lower()}_index" f"{attrs['Name'].lower()}_index"