From 154b070eaebca02d58a36931daa50548e793db31 Mon Sep 17 00:00:00 2001 From: Anders Melchiorsen Date: Tue, 26 Sep 2017 09:26:26 +0200 Subject: [PATCH] IMAP Unread sensor updated for async and push (#9562) * IMAP Unread sensor updated for async and push * Implement renames suggested in review * Use async_timeout * Keep push capability in a variable * Reword for Hound --- homeassistant/components/sensor/imap.py | 155 ++++++++++++++++++------ requirements_all.txt | 3 + 2 files changed, 118 insertions(+), 40 deletions(-) diff --git a/homeassistant/components/sensor/imap.py b/homeassistant/components/sensor/imap.py index 849f3fd8100..9d66537079f 100644 --- a/homeassistant/components/sensor/imap.py +++ b/homeassistant/components/sensor/imap.py @@ -5,20 +5,27 @@ For more details about this platform, please refer to the documentation at https://home-assistant.io/components/sensor.imap/ """ import logging +import asyncio +import async_timeout import voluptuous as vol from homeassistant.helpers.entity import Entity from homeassistant.components.sensor import PLATFORM_SCHEMA from homeassistant.const import ( - CONF_NAME, CONF_PORT, CONF_USERNAME, CONF_PASSWORD) + CONF_NAME, CONF_PORT, CONF_USERNAME, CONF_PASSWORD, + EVENT_HOMEASSISTANT_STOP) +from homeassistant.exceptions import PlatformNotReady import homeassistant.helpers.config_validation as cv _LOGGER = logging.getLogger(__name__) -CONF_SERVER = "server" +REQUIREMENTS = ['aioimaplib==0.7.12'] + +CONF_SERVER = 'server' DEFAULT_PORT = 993 + ICON = 'mdi:email-outline' PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({ @@ -30,17 +37,20 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({ }) -def setup_platform(hass, config, add_devices, discovery_info=None): - """Set up the IMAP platform.""" - sensor = ImapSensor( - config.get(CONF_NAME, None), config.get(CONF_USERNAME), - config.get(CONF_PASSWORD), config.get(CONF_SERVER), - config.get(CONF_PORT)) +@asyncio.coroutine +def async_setup_platform(hass, config, async_add_devices, discovery_info=None): + """Setup the IMAP platform.""" + sensor = ImapSensor(config.get(CONF_NAME), + config.get(CONF_USERNAME), + config.get(CONF_PASSWORD), + config.get(CONF_SERVER), + config.get(CONF_PORT)) - if sensor.connection: - add_devices([sensor], True) - else: - return False + if not (yield from sensor.connection()): + raise PlatformNotReady + + hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, sensor.shutdown()) + async_add_devices([sensor], True) class ImapSensor(Entity): @@ -54,45 +64,110 @@ class ImapSensor(Entity): self._server = server self._port = port self._unread_count = 0 - self.connection = self._login() + self._connection = None + self._does_push = None + self._idle_loop_task = None - def _login(self): - """Login and return an IMAP connection.""" - import imaplib - try: - connection = imaplib.IMAP4_SSL(self._server, self._port) - connection.login(self._user, self._password) - return connection - except imaplib.IMAP4.error: - _LOGGER.error("Failed to login to %s.", self._server) - return False + @asyncio.coroutine + def async_added_to_hass(self): + """Handle when an entity is about to be added to Home Assistant.""" + if not self.should_poll: + self._idle_loop_task = self.hass.loop.create_task(self.idle_loop()) @property def name(self): """Return the name of the sensor.""" return self._name + @property + def icon(self): + """Return the icon to use in the frontend.""" + return ICON + @property def state(self): """Return the number of unread emails.""" return self._unread_count - def update(self): - """Check the number of unread emails.""" - import imaplib - try: - self.connection.select() - self._unread_count = len(self.connection.search( - None, 'UnSeen UnDeleted')[1][0].split()) - except imaplib.IMAP4.error: - _LOGGER.info("Connection to %s lost, attempting to reconnect", - self._server) - try: - self.connection = self._login() - except imaplib.IMAP4.error: - _LOGGER.error("Failed to reconnect.") + @property + def available(self): + """Return the availability of the device.""" + return self._connection is not None @property - def icon(self): - """Return the icon to use in the frontend.""" - return ICON + def should_poll(self): + """Return if polling is needed.""" + return not self._does_push + + @asyncio.coroutine + def connection(self): + """Return a connection to the server, establishing it if necessary.""" + import aioimaplib + + if self._connection is None: + try: + self._connection = aioimaplib.IMAP4_SSL( + self._server, self._port) + yield from self._connection.wait_hello_from_server() + yield from self._connection.login(self._user, self._password) + yield from self._connection.select() + self._does_push = self._connection.has_capability('IDLE') + except (aioimaplib.AioImapException, asyncio.TimeoutError): + self._connection = None + + return self._connection + + @asyncio.coroutine + def idle_loop(self): + """Wait for data pushed from server.""" + import aioimaplib + + while True: + try: + if (yield from self.connection()): + yield from self.refresh_unread_count() + yield from self.async_update_ha_state() + + idle = yield from self._connection.idle_start() + yield from self._connection.wait_server_push() + self._connection.idle_done() + with async_timeout.timeout(10): + yield from idle + else: + yield from self.async_update_ha_state() + except (aioimaplib.AioImapException, asyncio.TimeoutError): + self.disconnected() + + @asyncio.coroutine + def async_update(self): + """Periodic polling of state.""" + import aioimaplib + + try: + if (yield from self.connection()): + yield from self.refresh_unread_count() + except (aioimaplib.AioImapException, asyncio.TimeoutError): + self.disconnected() + + @asyncio.coroutine + def refresh_unread_count(self): + """Check the number of unread emails.""" + if self._connection: + yield from self._connection.noop() + _, lines = yield from self._connection.search('UnSeen UnDeleted') + self._unread_count = len(lines[0].split()) + + def disconnected(self): + """Forget the connection after it was lost.""" + _LOGGER.warning("Lost %s (will attempt to reconnect)", self._server) + self._connection = None + + @asyncio.coroutine + def shutdown(self): + """Close resources.""" + if self._connection: + if self._connection.has_pending_idle(): + self._connection.idle_done() + yield from self._connection.logout() + if self._idle_loop_task: + self._idle_loop_task.cancel() diff --git a/requirements_all.txt b/requirements_all.txt index 659de15a860..022e88861af 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -57,6 +57,9 @@ aiodns==1.1.1 # homeassistant.components.http aiohttp_cors==0.5.3 +# homeassistant.components.sensor.imap +aioimaplib==0.7.12 + # homeassistant.components.light.lifx aiolifx==0.6.0