IMAP sensor async/await conversion (#12988)

homeassistant/components/sensor/imap.py

@@ -39,8 +39,10 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
 })
 
 
-@asyncio.coroutine
-def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
+async def async_setup_platform(hass,
+                               config,
+                               async_add_devices,
+                               discovery_info=None):
     """Set up the IMAP platform."""
     sensor = ImapSensor(config.get(CONF_NAME),
                         config.get(CONF_USERNAME),
@@ -48,8 +50,7 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
                         config.get(CONF_SERVER),
                         config.get(CONF_PORT),
                         config.get(CONF_FOLDER))
-
-    if not (yield from sensor.connection()):
+    if not await sensor.connection():
         raise PlatformNotReady
 
     hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, sensor.shutdown())
@@ -72,8 +73,7 @@ class ImapSensor(Entity):
         self._does_push = None
         self._idle_loop_task = None
 
-    @asyncio.coroutine
-    def async_added_to_hass(self):
+    async 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())
@@ -103,8 +103,7 @@ class ImapSensor(Entity):
         """Return if polling is needed."""
         return not self._does_push
 
-    @asyncio.coroutine
-    def connection(self):
+    async def connection(self):
         """Return a connection to the server, establishing it if necessary."""
         import aioimaplib
 
@@ -112,53 +111,50 @@ class ImapSensor(Entity):
             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._folder)
+                await self._connection.wait_hello_from_server()
+                await self._connection.login(self._user, self._password)
+                await self._connection.select(self._folder)
                 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):
+    async 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()
+                if await self.connection():
+                    await self.refresh_unread_count()
+                    await self.async_update_ha_state()
 
-                    idle = yield from self._connection.idle_start()
-                    yield from self._connection.wait_server_push()
+                    idle = await self._connection.idle_start()
+                    await self._connection.wait_server_push()
                     self._connection.idle_done()
                     with async_timeout.timeout(10):
-                        yield from idle
+                        await idle
                 else:
-                    yield from self.async_update_ha_state()
+                    await self.async_update_ha_state()
             except (aioimaplib.AioImapException, asyncio.TimeoutError):
                 self.disconnected()
 
-    @asyncio.coroutine
-    def async_update(self):
+    async def async_update(self):
         """Periodic polling of state."""
         import aioimaplib
 
         try:
-            if (yield from self.connection()):
-                yield from self.refresh_unread_count()
+            if await self.connection():
+                await self.refresh_unread_count()
         except (aioimaplib.AioImapException, asyncio.TimeoutError):
             self.disconnected()
 
-    @asyncio.coroutine
-    def refresh_unread_count(self):
+    async 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')
+            await self._connection.noop()
+            _, lines = await self._connection.search('UnSeen UnDeleted')
             self._unread_count = len(lines[0].split())
 
     def disconnected(self):
@@ -166,12 +162,11 @@ class ImapSensor(Entity):
         _LOGGER.warning("Lost %s (will attempt to reconnect)", self._server)
         self._connection = None
 
-    @asyncio.coroutine
-    def shutdown(self):
+    async def shutdown(self):
         """Close resources."""
         if self._connection:
             if self._connection.has_pending_idle():
                 self._connection.idle_done()
-            yield from self._connection.logout()
+            await self._connection.logout()
         if self._idle_loop_task:
             self._idle_loop_task.cancel()