diff --git a/Doc/library/imaplib.rst b/Doc/library/imaplib.rst --- a/Doc/library/imaplib.rst +++ b/Doc/library/imaplib.rst @@ -32,16 +32,29 @@ base class: .. class:: IMAP4(host='', port=IMAP4_PORT) This class implements the actual IMAP4 protocol. The connection is created and protocol version (IMAP4 or IMAP4rev1) is determined when the instance is 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 imap: + ... imap.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: .. exception:: IMAP4.error Exception raised on any errors. The reason for the exception is passed to the constructor as a string. diff --git a/Doc/whatsnew/3.5.rst b/Doc/whatsnew/3.5.rst --- a/Doc/whatsnew/3.5.rst +++ b/Doc/whatsnew/3.5.rst @@ -136,16 +136,24 @@ Improved Modules doctest ------- * :func:`doctest.DocTestSuite` returns an empty :class:`unittest.TestSuite` if *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 Serhiy Storchaka in + :issue:`4972`). + imghdr ------ * :func:`~imghdr.what` now recognizes the `OpenEXR `_ format (contributed by Martin vignali and Cladui Popa in :issue:`20295`). importlib --------- diff --git a/Lib/imaplib.py b/Lib/imaplib.py --- a/Lib/imaplib.py +++ b/Lib/imaplib.py @@ -233,16 +233,24 @@ class IMAP4: def __getattr__(self, attr): # Allow UPPERCASE variants of IMAP4 command methods. if attr in Commands: 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 def _create_socket(self): return socket.create_connection((self.host, self.port)) diff --git a/Lib/test/test_imaplib.py b/Lib/test/test_imaplib.py --- a/Lib/test/test_imaplib.py +++ b/Lib/test/test_imaplib.py @@ -96,16 +96,20 @@ else: class SimpleIMAPHandler(socketserver.StreamRequestHandler): timeout = 1 continuation = None capabilities = '' + def setup(self): + super().setup() + self.server.logged = None + def _send(self, message): if verbose: print("SENT: %r" % message.strip()) self.wfile.write(message) def _send_line(self, message): self._send(message + b'\r\n') def _send_textline(self, message): @@ -156,19 +160,24 @@ class SimpleIMAPHandler(socketserver.Str self._send_tagged(tag, 'BAD', cmd + ' unknown') def cmd_CAPABILITY(self, tag, args): caps = 'IMAP4rev1 ' + self.capabilities if self.capabilities else 'IMAP4rev1' self._send_textline('* CAPABILITY ' + caps) 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 BaseThreadedNetworkedTests(unittest.TestCase): def make_server(self, addr, hdlr): class MyServer(self.server_class): def handle_error(self, request, client_address): self.close_request(request) @@ -333,16 +342,42 @@ class BaseThreadedNetworkedTests(unittes def handle(self): # Send a very long response line self.wfile.write(b'* OK ' + imaplib._MAXLINE*b'x' + b'\r\n') with self.reaped_server(TooLongHandler) as server: 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) + class ThreadedNetworkedTests(BaseThreadedNetworkedTests): server_class = socketserver.TCPServer imap_class = imaplib.IMAP4 @unittest.skipUnless(ssl, "SSL not available")