diff --git a/homeassistant/components/iqvia/__init__.py b/homeassistant/components/iqvia/__init__.py index affe8622641..5a7fa682bc1 100644 --- a/homeassistant/components/iqvia/__init__.py +++ b/homeassistant/components/iqvia/__init__.py @@ -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 diff --git a/homeassistant/components/iqvia/const.py b/homeassistant/components/iqvia/const.py index 10b2ae30220..cbcda26982e 100644 --- a/homeassistant/components/iqvia/const.py +++ b/homeassistant/components/iqvia/const.py @@ -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"), -} diff --git a/homeassistant/components/iqvia/sensor.py b/homeassistant/components/iqvia/sensor.py index 068ba522e52..adf53a9cc05 100644 --- a/homeassistant/components/iqvia/sensor.py +++ b/homeassistant/components/iqvia/sensor.py @@ -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"