Added Search Configuration to IMAP Sensor

homeassistant/components/sensor/imap.py

@@ -24,6 +24,7 @@
 
 CONF_SERVER = 'server'
 CONF_FOLDER = 'folder'
+CONF_SEARCH = 'search'
 
 DEFAULT_PORT = 993
 
@@ -36,6 +37,7 @@
     vol.Required(CONF_SERVER): cv.string,
     vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
     vol.Optional(CONF_FOLDER, default='INBOX'): cv.string,
+    vol.Optional(CONF_SEARCH, default='UnSeen UnDeleted'): cv.string,
 })
 
 
@@ -49,7 +51,8 @@ async def async_setup_platform(hass,
                         config.get(CONF_PASSWORD),
                         config.get(CONF_SERVER),
                         config.get(CONF_PORT),
-                        config.get(CONF_FOLDER))
+                        config.get(CONF_FOLDER),
+                        config.get(CONF_SEARCH))
     if not await sensor.connection():
         raise PlatformNotReady
 
@@ -60,15 +63,16 @@ async def async_setup_platform(hass,
 class ImapSensor(Entity):
     """Representation of an IMAP sensor."""
 
-    def __init__(self, name, user, password, server, port, folder):
+    def __init__(self, name, user, password, server, port, folder, search):
         """Initialize the sensor."""
         self._name = name or user
         self._user = user
         self._password = password
         self._server = server
         self._port = port
         self._folder = folder
-        self._unread_count = 0
+        self._email_count = None
+        self._search = search
         self._connection = None
         self._does_push = None
         self._idle_loop_task = None
@@ -90,8 +94,8 @@ def icon(self):
 
     @property
     def state(self):
-        """Return the number of unread emails."""
-        return self._unread_count
+        """Return the number of emails found."""
+        return self._email_count
 
     @property
     def available(self):
@@ -127,7 +131,7 @@ async def idle_loop(self):
         while True:
             try:
                 if await self.connection():
-                    await self.refresh_unread_count()
+                    await self.refresh_email_count()
                     await self.async_update_ha_state()
 
                     idle = await self._connection.idle_start()
@@ -146,16 +150,22 @@ async def async_update(self):
 
         try:
             if await self.connection():
-                await self.refresh_unread_count()
+                await self.refresh_email_count()
         except (aioimaplib.AioImapException, asyncio.TimeoutError):
             self.disconnected()
 
-    async def refresh_unread_count(self):
-        """Check the number of unread emails."""
+    async def refresh_email_count(self):
+        """Check the number of found emails."""
         if self._connection:
             await self._connection.noop()
-            _, lines = await self._connection.search('UnSeen UnDeleted')
-            self._unread_count = len(lines[0].split())
+            result, lines = await self._connection.search(self._search)
+
+            if result == 'OK':
+                self._email_count = len(lines[0].split())
+            else:
+                _LOGGER.error("Can't parse IMAP server response to search "
+                              "'%s':  %s / %s",
+                              self._search, result, lines[0])
 
     def disconnected(self):
         """Forget the connection after it was lost."""