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

in a `with` statement.  Contributed by Giampaolo Rodola.
This commit is contained in:
Barry Warsaw 2011-03-15 15:04:44 -04:00
parent 5eb3591a41
commit 1f5c958721
4 changed files with 66 additions and 2 deletions

View File

@ -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])

View File

@ -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.

View File

@ -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.

View File

@ -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.