From 38684c366364fcd4071534a1de62f728407224d8 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Tue, 9 Sep 2014 19:07:49 +0300
Subject: [PATCH] =?UTF-8?q?imaplib.IMAP4=20now=20supports=20the=20context?=
 =?UTF-8?q?=20manager=20protocol.=20Original=20patch=20by=20Tarek=20Ziad?=
 =?UTF-8?q?=C3=A9.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 Doc/library/imaplib.rst  | 13 +++++++++++++
 Doc/whatsnew/3.5.rst     |  8 ++++++++
 Lib/imaplib.py           |  8 ++++++++
 Lib/test/test_imaplib.py | 35 +++++++++++++++++++++++++++++++++++
 Misc/NEWS                |  3 +++
 5 files changed, 67 insertions(+)

diff --git a/Doc/library/imaplib.rst b/Doc/library/imaplib.rst
index fa736fe3af4..f263bab9f02 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 43972b0f36a..38ea2f203a6 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 ad104fe76a7..27445dd6c51 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 910aa0e60a1..b58707359c6 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 157e72c52d6..ec19e7b2836 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