From 287d0a0a71ca0497e2d579391a565077b628d1ac Mon Sep 17 00:00:00 2001 From: Rami Mosleh Date: Tue, 9 Jun 2020 13:38:23 +0300 Subject: [PATCH] use HA DataUpdateCoordinator --- custom_components/trakt/__init__.py | 73 ++++++++++++++--------------- custom_components/trakt/sensor.py | 30 +++++++----- 2 files changed, 53 insertions(+), 50 deletions(-) diff --git a/custom_components/trakt/__init__.py b/custom_components/trakt/__init__.py index f9c3ff3..2442a4d 100644 --- a/custom_components/trakt/__init__.py +++ b/custom_components/trakt/__init__.py @@ -23,8 +23,7 @@ CONF_SCAN_INTERVAL, ) from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow -from homeassistant.helpers.dispatcher import async_dispatcher_send -from homeassistant.helpers.event import async_track_time_interval +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from trakt.calendar import MyShowCalendar from trakt.tv import TVShow @@ -32,13 +31,13 @@ from .const import ( CARD_DEFAULT, CONF_DAYS, - DATA_UPDATED, DEFAULT_DAYS, DEFAULT_SCAN_INTERVAL, DOMAIN, OAUTH2_AUTHORIZE, OAUTH2_TOKEN, ) +from homeassistant.exceptions import ConfigEntryNotReady _LOGGER = logging.getLogger(__name__) @@ -71,15 +70,16 @@ async def async_setup_entry(hass, entry) -> bool: session = config_entry_oauth2_flow.OAuth2Session(hass, entry, implementation) await session.async_ensure_token_valid() - trakt_data = Trakt_Data(hass, entry, session) - - if not await trakt_data.async_setup(): + coordinator = Trakt_Data(hass, entry, session) + if not await coordinator.async_setup(): return False - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = trakt_data + hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + hass.async_create_task( hass.config_entries.async_forward_entry_setup(entry, "sensor") ) + return True @@ -95,7 +95,7 @@ async def async_unload_entry(hass, entry) -> bool: return True -class Trakt_Data: +class Trakt_Data(DataUpdateCoordinator): """Represent Trakt data.""" def __init__( @@ -109,8 +109,19 @@ def __init__( self.config_entry = config_entry self.session = session self.unsub_timer = None - self.details = {} - self.calendar = None + self.calendar = {} + + super().__init__( + self.hass, + _LOGGER, + name=DOMAIN, + update_method=self.async_update, + update_interval=timedelta( + minutes=self.config_entry.options.get( + CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL + ) + ), + ) @property def days(self): @@ -135,15 +146,12 @@ async def async_update(self, *_): self.calendar = await self.hass.async_add_executor_job( MyShowCalendar, {CONF_DAYS: self.days} ) - except trakt.errors.OAuthException: - _LOGGER.error( - "Trakt api is unauthrized. Please remove the entry and reconfigure." - ) - return + except trakt.errors.TraktInternalException: + _LOGGER.error("Trakt api encountered an internal error.") + raise UpdateFailed if not self.calendar: _LOGGER.warning("Trakt upcoming calendar is empty") - return for show in self.calendar: if not show or show.show in self.exclude: @@ -203,10 +211,7 @@ async def async_update(self, *_): } card_json.append(card_item) - self.details = json.dumps(card_json) - - _LOGGER.debug("Trakt data updated") - async_dispatcher_send(self.hass, DATA_UPDATED) + return json.dumps(card_json) async def async_setup(self): """Set up Trakt Data.""" @@ -214,34 +219,26 @@ async def async_setup(self): trakt.core.CLIENT_ID = self.config_entry.data[CONF_CLIENT_ID] try: - await self.async_update() + await self.async_refresh() except trakt.errors.OAuthException: _LOGGER.error( "Trakt api is unauthrized. Please remove the entry and reconfigure." ) return False - await self.async_set_scan_interval(self.scan_interval) - self.config_entry.add_update_listener(self.async_options_updated) + if not self.last_update_success: + raise ConfigEntryNotReady + self.config_entry.add_update_listener(async_options_updated) return True - async def async_set_scan_interval(self, scan_interval): - """Update scan interval.""" - if self.unsub_timer is not None: - self.unsub_timer() - self.unsub_timer = async_track_time_interval( - self.hass, self.async_update, timedelta(minutes=scan_interval) - ) - await self.async_update() - - @staticmethod - async def async_options_updated(hass, entry): - """Triggered by config entry options updates.""" - await hass.data[DOMAIN][entry.entry_id].async_set_scan_interval( - entry.options[CONF_SCAN_INTERVAL] - ) +async def async_options_updated(hass, entry): + """Triggered by config entry options updates.""" + hass.data[DOMAIN][entry.entry_id].update_interval = timedelta( + minutes=entry.options[CONF_SCAN_INTERVAL] + ) + await hass.data[DOMAIN][entry.entry_id].async_request_refresh() def days_until(date): diff --git a/custom_components/trakt/sensor.py b/custom_components/trakt/sensor.py index fb95192..70b6caa 100644 --- a/custom_components/trakt/sensor.py +++ b/custom_components/trakt/sensor.py @@ -1,25 +1,24 @@ """Sensor platform for Trakt""" from homeassistant.const import ATTR_ATTRIBUTION, CONF_NAME, CONF_CLIENT_ID -from homeassistant.helpers.dispatcher import async_dispatcher_connect + from homeassistant.helpers.entity import Entity -from .const import ATTRIBUTION, DATA_UPDATED, DOMAIN +from .const import ATTRIBUTION, DOMAIN async def async_setup_entry(hass, config_entry, async_add_entities): """Set up device tracker for Mikrotik component.""" - tk_data = hass.data[DOMAIN][config_entry.entry_id] - - async_add_entities([TraktUpcomingCalendarSensor(tk_data)], True) + coordinator = hass.data[DOMAIN][config_entry.entry_id] + async_add_entities([TraktUpcomingCalendarSensor(coordinator)], True) class TraktUpcomingCalendarSensor(Entity): """Representation of a Trakt Upcoming Calendar sensor.""" - def __init__(self, tk_data): + def __init__(self, coordinator): """Initialize the sensor.""" - self.tk_data = tk_data - self._name = tk_data.config_entry.data[CONF_NAME] + self.coordinator = coordinator + self._name = coordinator.config_entry.data[CONF_NAME] @property def name(self): @@ -29,7 +28,7 @@ def name(self): @property def unique_id(self): """Return the unique id of the entity.""" - return self.tk_data.config_entry.data[CONF_CLIENT_ID] + return self.coordinator.config_entry.data[CONF_CLIENT_ID] @property def should_poll(self): @@ -39,7 +38,7 @@ def should_poll(self): @property def state(self): """Return the state of the sensor.""" - return len(self.tk_data.calendar) + return len(self.coordinator.calendar) if self.coordinator.calendar else 0 @property def icon(self): @@ -54,12 +53,19 @@ def unit_of_measurement(self): @property def device_state_attributes(self): """Return the state attributes of the sensor.""" - attributes = {"data": self.tk_data.details, ATTR_ATTRIBUTION: ATTRIBUTION} + attributes = { + "data": self.coordinator.data, + ATTR_ATTRIBUTION: ATTRIBUTION, + } return attributes async def async_added_to_hass(self): """Handle entity which will be added.""" self.async_on_remove( - async_dispatcher_connect(self.hass, DATA_UPDATED, self.async_write_ha_state) + self.coordinator.async_add_listener(self.async_write_ha_state) ) + + async def async_update(self): + """Request coordinator to update data.""" + await self.coordinator.async_request_refresh()