Fix imap re-connection

Imap re-connection is not working.

Connection exceptions (socket.error, socket.gaierror, IMAP4.abort) can
be raise during any imap operations, not only on connection
establishment.

Also the idle thread call start() on thread multiple times, but this is
not allowed by python. The thread is not daemonize, so the thread code
can be never runned if the python process is busy. And the module is
located in imaplib2.imaplib2.IMAP4 not imaplib2.IMAP4...

This change fixes all of that.

i3pystatus/mail/imap.py

@@ -1,60 +1,18 @@
 try:
-    from imaplib2 import IMAP4, IMAP4_SSL
+    from imaplib2.imaplib2 import IMAP4, IMAP4_SSL
     use_idle = True
-except:
+except ImportError:
     from imaplib import IMAP4, IMAP4_SSL
     use_idle = False
+import contextlib
+import time
 import socket
-from threading import *
+from threading import Thread
 
 from i3pystatus.mail import Backend
 
 
-# This is the threading object that does all the waiting on
+IMAP_EXCEPTIONS = (socket.error, socket.gaierror, IMAP4.abort, IMAP4.error)
-# the event
-class Idler(object):
-    def __init__(self, conn, callback, callback_reconnect):
-        self.thread = Thread(target=self.idle)
-        self.M = conn
-        self.callback_reconnect = callback_reconnect
-        self.event = Event()
-        self.callback = callback
-
-    def start(self):
-        self.thread.start()
-
-    def stop(self):
-        self.event.set()
-
-    def join(self):
-        self.thread.join()
-
-    def idle(self):
-
-        while True:
-            if self.event.isSet():
-                return
-
-            self.needsync = False
-
-            def callback(args):
-                if not self.event.isSet():
-                    self.needsync = True
-                    self.event.set()
-
-            try:
-                self.M.idle(callback=callback)
-                self.event.wait()
-
-                if self.needsync:
-                    self.event.clear()
-                    self.callback()
-            except:
-                break
-
-        self.M = self.callback_reconnect()
-        self.stop()
-        self.start()
 
 
 class IMAP(Backend):
@@ -84,47 +42,56 @@ class IMAP(Backend):
         if self.ssl:
             self.imap_class = IMAP4_SSL
 
-        self.conn = self.get_connection()
-
         if use_idle:
-            idler = Idler(self.conn, self.count_new_mail, self.get_connection)
+            self.thread = Thread(target=self._idle_thread)
-            idler.start()
+            self.daemon = True
+            self.thread.start()
 
-    def get_connection(self):
+    @contextlib.contextmanager
-        if self.connection:
+    def ensure_connection(self):
-            try:
+        try:
+            if self.connection:
                 self.connection.select(self.mailbox)
-            except socket.error:
+            if not self.connection:
-                # NOTE(sileht): retry just once if the connection have been
+                self.connection = self.imap_class(self.host, self.port)
-                # broken to ensure this is not a sporadic connection lost.
+                self.connection.login(self.username, self.password)
-                # Like wifi reconnect, sleep wake up
+                self.connection.select(self.mailbox)
-                try:
+            yield
-                    self.connection.logout()
+        except IMAP_EXCEPTIONS:
-                except socket.error:
+            # NOTE(sileht): retry just once if the connection have been
-                    pass
+            # broken to ensure this is not a sporadic connection lost.
-                self.connection = None
+            # Like wifi reconnect, sleep wake up
+            try:
+                self.connection.close()
+            except IMAP_EXCEPTIONS:
+                pass
+            try:
+                self.connection.logout()
+            except IMAP_EXCEPTIONS:
+                pass
+            # Wait a bit when disconnection occurs to not hog the cpu
+            time.sleep(1)
+            self.connection = None
 
-        if not self.connection:
+    def _idle_thread(self):
-            self.connection = self.imap_class(self.host, self.port)
+        # update mail count on startup
-            self.connection.login(self.username, self.password)
+        with self.ensure_connection():
-            self.connection.select(self.mailbox)
+            self.count_new_mail()
-
+        while True:
-        return self.connection
+            with self.ensure_connection():
+                # Block until new mails
+                self.connection.idle()
+                # Read how many
+                self.count_new_mail()
 
     def count_new_mail(self):
-        self.last = len(self.conn.search(None, "UnSeen")[1][0].split())
+        self.last = len(self.connection.search(None, "UnSeen")[1][0].split())
 
     @property
     def unread(self):
-        try:
+        if not use_idle:
-            conn = self.get_connection()
+            with self.ensure_connection():
-        except socket.gaierror:
-            pass
-        else:
-            self.last = len(conn.search(None, "UnSeen")[1][0].split())
-        finally:
-            if not use_idle:
                 self.count_new_mail()
-            return self.last
+        return self.last
 
 Backend = IMAP