- Issue #11289: `smtp.SMTP` class becomes a context manager so it can be used

  in a `with` statement.  Contributed by Giampaolo Rodola.

4 files changed, 66 insertions, 2 deletions

diff --git a/Doc/library/smtplib.rst b/Doc/library/smtplib.rst
index 531a64d..4805c8e 100644
--- a/Doc/library/smtplib.rst
+++ b/Doc/library/smtplib.rst
@@ -34,6 +34,20 @@ Protocol) and :rfc:`1869` (SMTP Service Extensions).
    For normal use, you should only require the initialization/connect,
    :meth:`sendmail`, and :meth:`quit` methods.  An example is included below.
 
+   The :class:`SMTP` class supports the :keyword:`with` statement.  When used
+   like this, the SMTP ``QUIT`` command is issued automatically when the
+   :keyword:`with` statement exits.  E.g.::
+
+    >>> from smtplib import SMTP
+    >>> with SMTP("domain.org") as smtp:
+    ...     smtp.noop()
+    ...
+    (250, b'Ok')
+    >>>
+
+   .. versionadded:: 3.3
+      Support for the :keyword:`with` statement was added.
+
 
 .. class:: SMTP_SSL(host='', port=0, local_hostname=None, keyfile=None, certfile=None[, timeout])
 
diff --git a/Lib/smtplib.py b/Lib/smtplib.py
index 14e6250..213138c 100755
--- a/Lib/smtplib.py
+++ b/Lib/smtplib.py
@@ -269,6 +269,19 @@ class SMTP:
                     pass
                 self.local_hostname = '[%s]' % addr
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        try:
+            code, message = self.docmd("QUIT")
+            if code != 221:
+                raise SMTPResponseException(code, message)
+        except SMTPServerDisconnected:
+            pass
+        finally:
+            self.close()
+
     def set_debuglevel(self, debuglevel):
         """Set the debug output level.
 
diff --git a/Lib/test/test_smtplib.py b/Lib/test/test_smtplib.py
index 4651f37..d973faa 100644
--- a/Lib/test/test_smtplib.py
+++ b/Lib/test/test_smtplib.py
@@ -424,6 +424,9 @@ sim_lists = {'list-1':['Mr.A@somewhere.com','Mrs.C@somewhereesle.com'],
 # Simulated SMTP channel & server
 class SimSMTPChannel(smtpd.SMTPChannel):
 
+    # For testing failures in QUIT when using the context manager API.
+    quit_response = None
+
     def __init__(self, extra_features, *args, **kw):
         self._extrafeatures = ''.join(
             [ "250-{0}\r\n".format(x) for x in extra_features ])
@@ -475,19 +478,31 @@ class SimSMTPChannel(smtpd.SMTPChannel):
         else:
             self.push('550 No access for you!')
 
+    def smtp_QUIT(self, arg):
+        # args is ignored
+        if self.quit_response is None:
+            super(SimSMTPChannel, self).smtp_QUIT(arg)
+        else:
+            self.push(self.quit_response)
+            self.close_when_done()
+
     def handle_error(self):
         raise
 
 
 class SimSMTPServer(smtpd.SMTPServer):
 
+    # For testing failures in QUIT when using the context manager API.
+    quit_response = None
+
     def __init__(self, *args, **kw):
         self._extra_features = []
         smtpd.SMTPServer.__init__(self, *args, **kw)
 
     def handle_accepted(self, conn, addr):
-        self._SMTPchannel = SimSMTPChannel(self._extra_features,
-                                           self, conn, addr)
+        self._SMTPchannel = SimSMTPChannel(
+            self._extra_features, self, conn, addr)
+        self._SMTPchannel.quit_response = self.quit_response
 
     def process_message(self, peer, mailfrom, rcpttos, data):
         pass
@@ -620,6 +635,25 @@ class SMTPSimTests(unittest.TestCase):
             self.assertIn(sim_auth_credentials['cram-md5'], str(err))
         smtp.close()
 
+    def test_with_statement(self):
+        with smtplib.SMTP(HOST, self.port) as smtp:
+            code, message = smtp.noop()
+            self.assertEqual(code, 250)
+        self.assertRaises(smtplib.SMTPServerDisconnected, smtp.send, b'foo')
+        with smtplib.SMTP(HOST, self.port) as smtp:
+            smtp.close()
+        self.assertRaises(smtplib.SMTPServerDisconnected, smtp.send, b'foo')
+
+    def test_with_statement_QUIT_failure(self):
+        self.serv.quit_response = '421 QUIT FAILED'
+        with self.assertRaises(smtplib.SMTPResponseException) as error:
+            with smtplib.SMTP(HOST, self.port) as smtp:
+                smtp.noop()
+        self.assertEqual(error.exception.smtp_code, 421)
+        self.assertEqual(error.exception.smtp_error, b'QUIT FAILED')
+        # We don't need to clean up self.serv.quit_response because a new
+        # server is always instantiated in the setUp().
+
     #TODO: add tests for correct AUTH method fallback now that the
     #test infrastructure can support it.
 
diff --git a/Misc/NEWS b/Misc/NEWS
index 834fa46..90e402b 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -68,6 +68,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #11289: `smtp.SMTP` class becomes a context manager so it can be used
+  in a `with` statement.  Contributed by Giampaolo Rodola.
+
 - Issue #11407: `TestCase.run` returns the result object used or created.
   Contributed by Janathan Hartley.