From 1b94ef69d89a7e4093093af59f623452857f6ed9 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Sat, 17 Oct 2020 10:16:41 -0600 Subject: [PATCH] Migrate IQVIA to DataUpdateCoordinator (#41970) * Migrate IQVIA to DataUpdateCoordinator * Linting * Code review * Better re-raise --- homeassistant/components/iqvia/__init__.py | 277 ++++++--------------- homeassistant/components/iqvia/const.py | 9 +- homeassistant/components/iqvia/sensor.py | 64 ++--- 3 files changed, 116 insertions(+), 234 deletions(-) diff --git a/homeassistant/components/iqvia/__init__.py b/homeassistant/components/iqvia/__init__.py index 171e7977e95d..9aa00ef3f3f7 100644 --- a/homeassistant/components/iqvia/__init__.py +++ b/homeassistant/components/iqvia/__init__.py @@ -1,238 +1,125 @@ """Support for IQVIA.""" import asyncio from datetime import timedelta -import logging from pyiqvia import Client -from pyiqvia.errors import InvalidZipError, IQVIAError +from pyiqvia.errors import IQVIAError from homeassistant.const import ATTR_ATTRIBUTION from homeassistant.core import callback from homeassistant.helpers import aiohttp_client -from homeassistant.helpers.dispatcher import ( - async_dispatcher_connect, - async_dispatcher_send, +from homeassistant.helpers.update_coordinator import ( + CoordinatorEntity, + DataUpdateCoordinator, + UpdateFailed, ) -from homeassistant.helpers.entity import Entity -from homeassistant.helpers.event import async_track_time_interval from .const import ( CONF_ZIP_CODE, - DATA_CLIENT, - DATA_LISTENER, + DATA_COORDINATOR, DOMAIN, - TOPIC_DATA_UPDATE, + LOGGER, TYPE_ALLERGY_FORECAST, TYPE_ALLERGY_INDEX, TYPE_ALLERGY_OUTLOOK, - TYPE_ALLERGY_TODAY, - TYPE_ALLERGY_TOMORROW, TYPE_ASTHMA_FORECAST, TYPE_ASTHMA_INDEX, - TYPE_ASTHMA_TODAY, - TYPE_ASTHMA_TOMORROW, TYPE_DISEASE_FORECAST, TYPE_DISEASE_INDEX, - TYPE_DISEASE_TODAY, ) -_LOGGER = logging.getLogger(__name__) - -API_CATEGORY_MAPPING = { - TYPE_ALLERGY_TODAY: TYPE_ALLERGY_INDEX, - TYPE_ALLERGY_TOMORROW: TYPE_ALLERGY_INDEX, - TYPE_ALLERGY_TOMORROW: TYPE_ALLERGY_INDEX, - TYPE_ASTHMA_TODAY: TYPE_ASTHMA_INDEX, - TYPE_ASTHMA_TOMORROW: TYPE_ALLERGY_INDEX, - TYPE_DISEASE_TODAY: TYPE_DISEASE_INDEX, -} - -DATA_CONFIG = "config" - DEFAULT_ATTRIBUTION = "Data provided by IQVIA™" DEFAULT_SCAN_INTERVAL = timedelta(minutes=30) - -@callback -def async_get_api_category(sensor_type): - """Return the API category that a particular sensor type should use.""" - return API_CATEGORY_MAPPING.get(sensor_type, sensor_type) +PLATFORMS = ["sensor"] async def async_setup(hass, config): """Set up the IQVIA component.""" - hass.data[DOMAIN] = {} - hass.data[DOMAIN][DATA_CLIENT] = {} - hass.data[DOMAIN][DATA_LISTENER] = {} - + hass.data[DOMAIN] = {DATA_COORDINATOR: {}} return True -async def async_setup_entry(hass, config_entry): +async def async_setup_entry(hass, entry): """Set up IQVIA as config entry.""" - websession = aiohttp_client.async_get_clientsession(hass) + hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] = {} - if not config_entry.unique_id: + if not entry.unique_id: # If the config entry doesn't already have a unique ID, set one: hass.config_entries.async_update_entry( - config_entry, **{"unique_id": config_entry.data[CONF_ZIP_CODE]} + entry, **{"unique_id": entry.data[CONF_ZIP_CODE]} ) - iqvia = IQVIAData(hass, Client(config_entry.data[CONF_ZIP_CODE], websession)) + websession = aiohttp_client.async_get_clientsession(hass) + client = Client(entry.data[CONF_ZIP_CODE], websession) - try: - await iqvia.async_update() - except InvalidZipError: - _LOGGER.error("Invalid ZIP code provided: %s", config_entry.data[CONF_ZIP_CODE]) - return False + async def async_get_data_from_api(api_coro): + """Get data from a particular API coroutine.""" + try: + return await api_coro() + except IQVIAError as err: + raise UpdateFailed from err - hass.data[DOMAIN][DATA_CLIENT][config_entry.entry_id] = iqvia + init_data_update_tasks = [] + for sensor_type, api_coro in [ + (TYPE_ALLERGY_FORECAST, client.allergens.extended), + (TYPE_ALLERGY_INDEX, client.allergens.current), + (TYPE_ALLERGY_OUTLOOK, client.allergens.outlook), + (TYPE_ASTHMA_FORECAST, client.asthma.extended), + (TYPE_ASTHMA_INDEX, client.asthma.current), + (TYPE_DISEASE_FORECAST, client.disease.extended), + (TYPE_DISEASE_INDEX, client.disease.current), + ]: + coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][ + sensor_type + ] = DataUpdateCoordinator( + hass, + LOGGER, + name=f"{entry.data[CONF_ZIP_CODE]} {sensor_type}", + update_interval=DEFAULT_SCAN_INTERVAL, + update_method=lambda coro=api_coro: async_get_data_from_api(coro), + ) + init_data_update_tasks.append(coordinator.async_refresh()) - hass.async_create_task( - hass.config_entries.async_forward_entry_setup(config_entry, "sensor") + await asyncio.gather(*init_data_update_tasks) + + for component in PLATFORMS: + hass.async_create_task( + hass.config_entries.async_forward_entry_setup(entry, component) + ) + + return True + + +async def async_unload_entry(hass, entry): + """Unload an OpenUV config entry.""" + unload_ok = all( + await asyncio.gather( + *[ + hass.config_entries.async_forward_entry_unload(entry, component) + for component in PLATFORMS + ] + ) ) - return True + if unload_ok: + hass.data[DOMAIN][DATA_COORDINATOR].pop(entry.entry_id) + + return unload_ok -async def async_unload_entry(hass, config_entry): - """Unload an OpenUV config entry.""" - hass.data[DOMAIN][DATA_CLIENT].pop(config_entry.entry_id) - - remove_listener = hass.data[DOMAIN][DATA_LISTENER].pop(config_entry.entry_id) - remove_listener() - - await hass.config_entries.async_forward_entry_unload(config_entry, "sensor") - - return True - - -class IQVIAData: - """Define a data object to retrieve info from IQVIA.""" - - def __init__(self, hass, client): - """Initialize.""" - self._async_cancel_time_interval_listener = None - self._client = client - self._hass = hass - self.data = {} - self.zip_code = client.zip_code - - self._api_coros = { - TYPE_ALLERGY_FORECAST: client.allergens.extended, - TYPE_ALLERGY_INDEX: client.allergens.current, - TYPE_ALLERGY_OUTLOOK: client.allergens.outlook, - TYPE_ASTHMA_FORECAST: client.asthma.extended, - TYPE_ASTHMA_INDEX: client.asthma.current, - TYPE_DISEASE_FORECAST: client.disease.extended, - TYPE_DISEASE_INDEX: client.disease.current, - } - self._api_category_count = { - TYPE_ALLERGY_FORECAST: 0, - TYPE_ALLERGY_INDEX: 0, - TYPE_ALLERGY_OUTLOOK: 0, - TYPE_ASTHMA_FORECAST: 0, - TYPE_ASTHMA_INDEX: 0, - TYPE_DISEASE_FORECAST: 0, - TYPE_DISEASE_INDEX: 0, - } - self._api_category_locks = { - TYPE_ALLERGY_FORECAST: asyncio.Lock(), - TYPE_ALLERGY_INDEX: asyncio.Lock(), - TYPE_ALLERGY_OUTLOOK: asyncio.Lock(), - TYPE_ASTHMA_FORECAST: asyncio.Lock(), - TYPE_ASTHMA_INDEX: asyncio.Lock(), - TYPE_DISEASE_FORECAST: asyncio.Lock(), - TYPE_DISEASE_INDEX: asyncio.Lock(), - } - - async def _async_get_data_from_api(self, api_category): - """Update and save data for a particular API category.""" - if self._api_category_count[api_category] == 0: - return - - try: - self.data[api_category] = await self._api_coros[api_category]() - except IQVIAError as err: - _LOGGER.error("Unable to get %s data: %s", api_category, err) - self.data[api_category] = None - - async def _async_update_listener_action(self, now): - """Define an async_track_time_interval action to update data.""" - await self.async_update() - - @callback - def async_deregister_api_interest(self, sensor_type): - """Decrement the number of entities with data needs from an API category.""" - # If this deregistration should leave us with no registration at all, remove the - # time interval: - if sum(self._api_category_count.values()) == 0: - if self._async_cancel_time_interval_listener: - self._async_cancel_time_interval_listener() - self._async_cancel_time_interval_listener = None - return - - api_category = async_get_api_category(sensor_type) - self._api_category_count[api_category] -= 1 - - async def async_register_api_interest(self, sensor_type): - """Increment the number of entities with data needs from an API category.""" - # If this is the first registration we have, start a time interval: - if not self._async_cancel_time_interval_listener: - self._async_cancel_time_interval_listener = async_track_time_interval( - self._hass, - self._async_update_listener_action, - DEFAULT_SCAN_INTERVAL, - ) - - api_category = async_get_api_category(sensor_type) - self._api_category_count[api_category] += 1 - - # If a sensor registers interest in a particular API call and the data doesn't - # exist for it yet, make the API call and grab the data: - async with self._api_category_locks[api_category]: - if api_category not in self.data: - await self._async_get_data_from_api(api_category) - - async def async_update(self): - """Update IQVIA data.""" - tasks = [ - self._async_get_data_from_api(api_category) - for api_category in self._api_coros - ] - - await asyncio.gather(*tasks) - - _LOGGER.debug("Received new data") - async_dispatcher_send(self._hass, TOPIC_DATA_UPDATE) - - -class IQVIAEntity(Entity): +class IQVIAEntity(CoordinatorEntity): """Define a base IQVIA entity.""" - def __init__(self, iqvia, sensor_type, name, icon, zip_code): - """Initialize the sensor.""" + def __init__(self, coordinator, entry, sensor_type, name, icon): + """Initialize.""" + super().__init__(coordinator) self._attrs = {ATTR_ATTRIBUTION: DEFAULT_ATTRIBUTION} + self._entry = entry self._icon = icon - self._iqvia = iqvia self._name = name self._state = None self._type = sensor_type - self._zip_code = zip_code - - @property - def available(self): - """Return True if entity is available.""" - if self._type in (TYPE_ALLERGY_TODAY, TYPE_ALLERGY_TOMORROW): - return self._iqvia.data.get(TYPE_ALLERGY_INDEX) is not None - - if self._type in (TYPE_ASTHMA_TODAY, TYPE_ASTHMA_TOMORROW): - return self._iqvia.data.get(TYPE_ASTHMA_INDEX) is not None - - if self._type == TYPE_DISEASE_TODAY: - return self._iqvia.data.get(TYPE_DISEASE_INDEX) is not None - - return self._iqvia.data.get(self._type) is not None @property def device_state_attributes(self): @@ -257,7 +144,7 @@ class IQVIAEntity(Entity): @property def unique_id(self): """Return a unique, Home Assistant friendly identifier for this entity.""" - return f"{self._zip_code}_{self._type}" + return f"{self._entry.data[CONF_ZIP_CODE]}_{self._type}" @property def unit_of_measurement(self): @@ -273,27 +160,17 @@ class IQVIAEntity(Entity): self.update_from_latest_data() self.async_write_ha_state() - self.async_on_remove( - async_dispatcher_connect(self.hass, TOPIC_DATA_UPDATE, update) - ) + self.async_on_remove(self.coordinator.async_add_listener(update)) - await self._iqvia.async_register_api_interest(self._type) if self._type == TYPE_ALLERGY_FORECAST: - # Entities that express interest in allergy forecast data should also - # express interest in allergy outlook data: - await self._iqvia.async_register_api_interest(TYPE_ALLERGY_OUTLOOK) + outlook_coordinator = self.hass.data[DOMAIN][DATA_COORDINATOR][ + self._entry.entry_id + ][TYPE_ALLERGY_OUTLOOK] + self.async_on_remove(outlook_coordinator.async_add_listener(update)) self.update_from_latest_data() - async def async_will_remove_from_hass(self): - """Disconnect dispatcher listener when removed.""" - self._iqvia.async_deregister_api_interest(self._type) - if self._type == TYPE_ALLERGY_FORECAST: - # Entities that lose interest in allergy forecast data should also lose - # interest in allergy outlook data: - self._iqvia.async_deregister_api_interest(TYPE_ALLERGY_OUTLOOK) - @callback def update_from_latest_data(self): - """Update the entity's state.""" - raise NotImplementedError() + """Update the entity from the latest data.""" + raise NotImplementedError diff --git a/homeassistant/components/iqvia/const.py b/homeassistant/components/iqvia/const.py index 95b03485597f..10b2ae302202 100644 --- a/homeassistant/components/iqvia/const.py +++ b/homeassistant/components/iqvia/const.py @@ -1,12 +1,13 @@ """Define IQVIA constants.""" +import logging + +LOGGER = logging.getLogger(__package__) + DOMAIN = "iqvia" CONF_ZIP_CODE = "zip_code" -DATA_CLIENT = "client" -DATA_LISTENER = "listener" - -TOPIC_DATA_UPDATE = f"{DOMAIN}_data_update" +DATA_COORDINATOR = "coordinator" TYPE_ALLERGY_FORECAST = "allergy_average_forecasted" TYPE_ALLERGY_INDEX = "allergy_index" diff --git a/homeassistant/components/iqvia/sensor.py b/homeassistant/components/iqvia/sensor.py index 9f9c2d534fc1..e53ee96b1c40 100644 --- a/homeassistant/components/iqvia/sensor.py +++ b/homeassistant/components/iqvia/sensor.py @@ -3,9 +3,14 @@ from statistics import mean import numpy as np -from homeassistant.components.iqvia import ( - DATA_CLIENT, +from homeassistant.const import ATTR_STATE +from homeassistant.core import callback + +from . import IQVIAEntity +from .const import ( + DATA_COORDINATOR, DOMAIN, + SENSORS, TYPE_ALLERGY_FORECAST, TYPE_ALLERGY_INDEX, TYPE_ALLERGY_OUTLOOK, @@ -18,12 +23,7 @@ from homeassistant.components.iqvia import ( TYPE_DISEASE_FORECAST, TYPE_DISEASE_INDEX, TYPE_DISEASE_TODAY, - IQVIAEntity, ) -from homeassistant.const import ATTR_STATE -from homeassistant.core import callback - -from .const import SENSORS ATTR_ALLERGEN_AMOUNT = "allergen_amount" ATTR_ALLERGEN_GENUS = "allergen_genus" @@ -36,6 +36,15 @@ ATTR_SEASON = "season" ATTR_TREND = "trend" ATTR_ZIP_CODE = "zip_code" +API_CATEGORY_MAPPING = { + TYPE_ALLERGY_TODAY: TYPE_ALLERGY_INDEX, + TYPE_ALLERGY_TOMORROW: TYPE_ALLERGY_INDEX, + TYPE_ALLERGY_TOMORROW: TYPE_ALLERGY_INDEX, + TYPE_ASTHMA_TODAY: TYPE_ASTHMA_INDEX, + TYPE_ASTHMA_TOMORROW: TYPE_ASTHMA_INDEX, + TYPE_DISEASE_TODAY: TYPE_DISEASE_INDEX, +} + RATING_MAPPING = [ {"label": "Low", "minimum": 0.0, "maximum": 2.4}, {"label": "Low/Medium", "minimum": 2.5, "maximum": 4.8}, @@ -51,8 +60,6 @@ TREND_SUBSIDING = "Subsiding" async def async_setup_entry(hass, entry, async_add_entities): """Set up IQVIA sensors based on a config entry.""" - iqvia = hass.data[DOMAIN][DATA_CLIENT][entry.entry_id] - sensor_class_mapping = { TYPE_ALLERGY_FORECAST: ForecastSensor, TYPE_ALLERGY_TODAY: IndexSensor, @@ -64,14 +71,15 @@ async def async_setup_entry(hass, entry, async_add_entities): TYPE_DISEASE_TODAY: IndexSensor, } - async_add_entities( - [ - sensor_class_mapping[sensor_type]( - iqvia, sensor_type, name, icon, iqvia.zip_code - ) - for sensor_type, (name, icon) in SENSORS.items() - ] - ) + 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)) + + async_add_entities(sensors) def calculate_trend(indices): @@ -96,10 +104,7 @@ class ForecastSensor(IQVIAEntity): @callback def update_from_latest_data(self): """Update the sensor.""" - if not self._iqvia.data.get(self._type): - return - - data = self._iqvia.data[self._type].get("Location") + data = self.coordinator.data.get("Location") if not data or not data.get("periods"): return @@ -122,9 +127,11 @@ class ForecastSensor(IQVIAEntity): ) if self._type == TYPE_ALLERGY_FORECAST: - outlook = self._iqvia.data[TYPE_ALLERGY_OUTLOOK] - self._attrs[ATTR_OUTLOOK] = outlook.get("Outlook") - self._attrs[ATTR_SEASON] = outlook.get("Season") + outlook_coordinator = self.hass.data[DOMAIN][DATA_COORDINATOR][ + self._entry.entry_id + ][TYPE_ALLERGY_OUTLOOK] + self._attrs[ATTR_OUTLOOK] = outlook_coordinator.data.get("Outlook") + self._attrs[ATTR_SEASON] = outlook_coordinator.data.get("Season") self._state = average @@ -135,16 +142,13 @@ class IndexSensor(IQVIAEntity): @callback def update_from_latest_data(self): """Update the sensor.""" - if not self._iqvia.data: - return - try: if self._type in (TYPE_ALLERGY_TODAY, TYPE_ALLERGY_TOMORROW): - data = self._iqvia.data[TYPE_ALLERGY_INDEX].get("Location") + data = self.coordinator.data.get("Location") elif self._type in (TYPE_ASTHMA_TODAY, TYPE_ASTHMA_TOMORROW): - data = self._iqvia.data[TYPE_ASTHMA_INDEX].get("Location") + data = self.coordinator.data.get("Location") elif self._type == TYPE_DISEASE_TODAY: - data = self._iqvia.data[TYPE_DISEASE_INDEX].get("Location") + data = self.coordinator.data.get("Location") except KeyError: return