imaplib.IMAP4 now supports the context manager protocol.

Original patch by Tarek Ziadé.

5 files changed, 67 insertions, 0 deletions

diff --git a/Doc/library/imaplib.rst b/Doc/library/imaplib.rst
index fa736fe..f263bab 100644
--- a/Doc/library/imaplib.rst
+++ b/Doc/library/imaplib.rst
@@ -37,6 +37,19 @@ base class:
    initialized. If *host* is not specified, ``''`` (the local host) is used. If
    *port* is omitted, the standard IMAP4 port (143) is used.
 
+   The :class:`IMAP4` class supports the :keyword:`with` statement.  When used
+   like this, the IMAP4 ``LOGOUT`` command is issued automatically when the
+   :keyword:`with` statement exits.  E.g.::
+
+    >>> from imaplib import IMAP4
+    >>> with IMAP4("domain.org") as M:
+    ...     M.noop()
+    ...
+    ('OK', [b'Nothing Accomplished. d25if65hy903weo.87'])
+
+   .. versionchanged:: 3.5
+      Support for the :keyword:`with` statement was added.
+
 Three exceptions are defined as attributes of the :class:`IMAP4` class:
 
 
diff --git a/Doc/whatsnew/3.5.rst b/Doc/whatsnew/3.5.rst
index 43972b0..38ea2f2 100644
--- a/Doc/whatsnew/3.5.rst
+++ b/Doc/whatsnew/3.5.rst
@@ -141,6 +141,14 @@ doctest
   *module* contains no docstrings instead of raising :exc:`ValueError`
   (contributed by Glenn Jones in :issue:`15916`).
 
+imaplib
+-------
+
+* :class:`IMAP4` now supports the context management protocol.  When used in a
+  :keyword:`with` statement, the IMAP4 ``LOGOUT`` command will be called
+  automatically at the end of the block.  (Contributed by Tarek Ziadé and
+  Serhiy Storchaka in :issue:`4972`).
+
 imghdr
 ------
 
diff --git a/Lib/imaplib.py b/Lib/imaplib.py
index ad104fe..27445dd 100644
--- a/Lib/imaplib.py
+++ b/Lib/imaplib.py
@@ -238,6 +238,14 @@ class IMAP4:
             return getattr(self, attr.lower())
         raise AttributeError("Unknown IMAP4 command: '%s'" % attr)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        try:
+            self.logout()
+        except OSError:
+            pass
 
 
     #       Overridable methods
diff --git a/Lib/test/test_imaplib.py b/Lib/test/test_imaplib.py
index 910aa0e..b587073 100644
--- a/Lib/test/test_imaplib.py
+++ b/Lib/test/test_imaplib.py
@@ -98,6 +98,10 @@ class SimpleIMAPHandler(socketserver.StreamRequestHandler):
     continuation = None
     capabilities = ''
 
+    def setup(self):
+        super().setup()
+        self.server.logged = None
+
     def _send(self, message):
         if verbose:
             print("SENT: %r" % message.strip())
@@ -162,9 +166,14 @@ class SimpleIMAPHandler(socketserver.StreamRequestHandler):
         self._send_tagged(tag, 'OK', 'CAPABILITY completed')
 
     def cmd_LOGOUT(self, tag, args):
+        self.server.logged = None
         self._send_textline('* BYE IMAP4ref1 Server logging out')
         self._send_tagged(tag, 'OK', 'LOGOUT completed')
 
+    def cmd_LOGIN(self, tag, args):
+        self.server.logged = args[0]
+        self._send_tagged(tag, 'OK', 'LOGIN completed')
+
 
 class ThreadedNetworkedTests(unittest.TestCase):
     server_class = socketserver.TCPServer
@@ -345,6 +354,32 @@ class ThreadedNetworkedTests(unittest.TestCase):
             self.assertRaises(imaplib.IMAP4.error,
                               self.imap_class, *server.server_address)
 
+    @reap_threads
+    def test_simple_with_statement(self):
+        # simplest call
+        with self.reaped_server(SimpleIMAPHandler) as server:
+            with self.imap_class(*server.server_address):
+                pass
+
+    @reap_threads
+    def test_with_statement(self):
+        with self.reaped_server(SimpleIMAPHandler) as server:
+            with self.imap_class(*server.server_address) as imap:
+                imap.login('user', 'pass')
+                self.assertEqual(server.logged, 'user')
+            self.assertIsNone(server.logged)
+
+    @reap_threads
+    def test_with_statement_logout(self):
+        # what happens if already logout in the block?
+        with self.reaped_server(SimpleIMAPHandler) as server:
+            with self.imap_class(*server.server_address) as imap:
+                imap.login('user', 'pass')
+                self.assertEqual(server.logged, 'user')
+                imap.logout()
+                self.assertIsNone(server.logged)
+            self.assertIsNone(server.logged)
+
 
 @unittest.skipUnless(ssl, "SSL not available")
 class ThreadedNetworkedTestsSSL(ThreadedNetworkedTests):
diff --git a/Misc/NEWS b/Misc/NEWS
index 157e72c..ec19e7b 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -132,6 +132,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #12410: imaplib.IMAP4 now supports the context manager protocol.
+  Original patch by Tarek Ziadé.
+
 - Issue #16662: load_tests() is now unconditionally run when it is present in
   a package's __init__.py.  TestLoader.loadTestsFromModule() still accepts
   use_load_tests, but it is deprecated and ignored.  A new keyword-only