bpo-38615: Add timeout parameter for IMAP4 and IMAP4_SSL constructor (GH-17203)

imaplib.IMAP4 and imaplib.IMAP4_SSL now have an 
optional *timeout* parameter for their constructors.
Also, the imaplib.IMAP4.open() method now has an optional *timeout* parameter
with this change. The overridden methods of imaplib.IMAP4_SSL and
imaplib.IMAP4_stream were applied to this change.

5 files changed, 90 insertions, 25 deletions

diff --git a/Doc/library/imaplib.rst b/Doc/library/imaplib.rst
index df63d82..5b8ca7c 100644
--- a/Doc/library/imaplib.rst
+++ b/Doc/library/imaplib.rst
@@ -30,12 +30,14 @@ Three classes are provided by the :mod:`imaplib` module, :class:`IMAP4` is the
 base class:
 
 
-.. class:: IMAP4(host='', port=IMAP4_PORT)
+.. class:: IMAP4(host='', port=IMAP4_PORT, timeout=None)
 
    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.
+   *port* is omitted, the standard IMAP4 port (143) is used. The optional *timeout*
+   parameter specifies a timeout in seconds for the connection attempt.
+   If timeout is not given or is None, the global default socket timeout is used.
 
    The :class:`IMAP4` class supports the :keyword:`with` statement.  When used
    like this, the IMAP4 ``LOGOUT`` command is issued automatically when the
@@ -50,6 +52,9 @@ base class:
    .. versionchanged:: 3.5
       Support for the :keyword:`with` statement was added.
 
+   .. versionchanged:: 3.9
+      The optional *timeout* parameter was added.
+
 Three exceptions are defined as attributes of the :class:`IMAP4` class:
 
 
@@ -78,7 +83,7 @@ There's also a subclass for secure connections:
 
 
 .. class:: IMAP4_SSL(host='', port=IMAP4_SSL_PORT, keyfile=None, \
-                     certfile=None, ssl_context=None)
+                     certfile=None, ssl_context=None, timeout=None)
 
    This is a subclass derived from :class:`IMAP4` that connects over an SSL
    encrypted socket (to use this class you need a socket module that was compiled
@@ -95,8 +100,12 @@ There's also a subclass for secure connections:
    mutually exclusive with *ssl_context*, a :class:`ValueError` is raised
    if *keyfile*/*certfile* is provided along with *ssl_context*.
 
+   The optional *timeout* parameter specifies a timeout in seconds for the
+   connection attempt. If timeout is not given or is None, the global default
+   socket timeout is used.
+
    .. versionchanged:: 3.3
-      *ssl_context* parameter added.
+      *ssl_context* parameter was added.
 
    .. versionchanged:: 3.4
       The class now supports hostname check with
@@ -110,6 +119,8 @@ There's also a subclass for secure connections:
        :func:`ssl.create_default_context` select the system's trusted CA
        certificates for you.
 
+   .. versionchanged:: 3.9
+      The optional *timeout* parameter was added.
 
 The second subclass allows for connections created by a child process:
 
@@ -353,16 +364,22 @@ An :class:`IMAP4` instance has the following methods:
    Send ``NOOP`` to server.
 
 
-.. method:: IMAP4.open(host, port)
+.. method:: IMAP4.open(host, port, timeout=None)
 
-   Opens socket to *port* at *host*.  This method is implicitly called by
-   the :class:`IMAP4` constructor.  The connection objects established by this
-   method will be used in the :meth:`IMAP4.read`, :meth:`IMAP4.readline`,
-   :meth:`IMAP4.send`, and :meth:`IMAP4.shutdown` methods.  You may override
-   this method.
+   Opens socket to *port* at *host*. The optional *timeout* parameter
+   specifies a timeout in seconds for the connection attempt.
+   If timeout is not given or is None, the global default socket timeout
+   is used. Also note that if the *timeout* parameter is set to be zero,
+   it will raise a :class:`ValueError` to reject creating a non-blocking socket.
+   This method is implicitly called by the :class:`IMAP4` constructor.
+   The connection objects established by this method will be used in
+   the :meth:`IMAP4.read`, :meth:`IMAP4.readline`, :meth:`IMAP4.send`,
+   and :meth:`IMAP4.shutdown` methods. You may override this method.
 
    .. audit-event:: imaplib.open self,host,port imaplib.IMAP4.open
 
+   .. versionchanged:: 3.9
+      The *timeout* parameter was added.
 
 .. method:: IMAP4.partial(message_num, message_part, start, length)
 
diff --git a/Doc/whatsnew/3.9.rst b/Doc/whatsnew/3.9.rst
index 46774c2..ea6d8f5 100644
--- a/Doc/whatsnew/3.9.rst
+++ b/Doc/whatsnew/3.9.rst
@@ -167,6 +167,16 @@ When the garbage collector makes a collection in which some objects resurrect
 been executed), do not block the collection of all objects that are still
 unreachable. (Contributed by Pablo Galindo and Tim Peters in :issue:`38379`.)
 
+imaplib
+-------
+
+:class:`~imaplib.IMAP4` and :class:`~imaplib.IMAP4_SSL` now have
+an optional *timeout* parameter for their constructors.
+Also, the :meth:`~imaplib.IMAP4.open` method now has an optional *timeout* parameter
+with this change. The overridden methods of :class:`~imaplib.IMAP4_SSL` and
+:class:`~imaplib.IMAP4_stream` were applied to this change.
+(Contributed by Dong-hee Na in :issue:`38615`.)
+
 os
 --
 
diff --git a/Lib/imaplib.py b/Lib/imaplib.py
index a4f4993..abfdd73 100644
--- a/Lib/imaplib.py
+++ b/Lib/imaplib.py
@@ -135,10 +135,13 @@ class IMAP4:
 
     r"""IMAP4 client class.
 
-    Instantiate with: IMAP4([host[, port]])
+    Instantiate with: IMAP4([host[, port[, timeout=None]]])
 
             host - host's name (default: localhost);
             port - port number (default: standard IMAP4 port).
+            timeout - socket timeout (default: None)
+                      If timeout is not given or is None,
+                      the global default socket timeout is used
 
     All IMAP4rev1 commands are supported by methods of the same
     name (in lower-case).
@@ -181,7 +184,7 @@ class IMAP4:
     class abort(error): pass        # Service errors - close and retry
     class readonly(abort): pass     # Mailbox status changed to READ-ONLY
 
-    def __init__(self, host='', port=IMAP4_PORT):
+    def __init__(self, host='', port=IMAP4_PORT, timeout=None):
         self.debug = Debug
         self.state = 'LOGOUT'
         self.literal = None             # A literal argument to a command
@@ -195,7 +198,7 @@ class IMAP4:
 
         # Open socket to server.
 
-        self.open(host, port)
+        self.open(host, port, timeout)
 
         try:
             self._connect()
@@ -284,15 +287,20 @@ class IMAP4:
     #       Overridable methods
 
 
-    def _create_socket(self):
+    def _create_socket(self, timeout):
         # Default value of IMAP4.host is '', but socket.getaddrinfo()
         # (which is used by socket.create_connection()) expects None
         # as a default value for host.
+        if timeout is not None and not timeout:
+            raise ValueError('Non-blocking socket (timeout=0) is not supported')
         host = None if not self.host else self.host
         sys.audit("imaplib.open", self, self.host, self.port)
-        return socket.create_connection((host, self.port))
+        address = (host, self.port)
+        if timeout is not None:
+            return socket.create_connection(address, timeout)
+        return socket.create_connection(address)
 
-    def open(self, host = '', port = IMAP4_PORT):
+    def open(self, host='', port=IMAP4_PORT, timeout=None):
         """Setup connection to remote server on "host:port"
             (default: localhost:standard IMAP4 port).
         This connection will be used by the routines:
@@ -300,7 +308,7 @@ class IMAP4:
         """
         self.host = host
         self.port = port
-        self.sock = self._create_socket()
+        self.sock = self._create_socket(timeout)
         self.file = self.sock.makefile('rb')
 
 
@@ -1261,7 +1269,7 @@ if HAVE_SSL:
 
         """IMAP4 client class over SSL connection
 
-        Instantiate with: IMAP4_SSL([host[, port[, keyfile[, certfile[, ssl_context]]]]])
+        Instantiate with: IMAP4_SSL([host[, port[, keyfile[, certfile[, ssl_context[, timeout=None]]]]]])
 
                 host - host's name (default: localhost);
                 port - port number (default: standard IMAP4 SSL port);
@@ -1271,13 +1279,15 @@ if HAVE_SSL:
                               and private key (default: None)
                 Note: if ssl_context is provided, then parameters keyfile or
                 certfile should not be set otherwise ValueError is raised.
+                timeout - socket timeout (default: None) If timeout is not given or is None,
+                          the global default socket timeout is used
 
         for more documentation see the docstring of the parent class IMAP4.
         """
 
 
         def __init__(self, host='', port=IMAP4_SSL_PORT, keyfile=None,
-                     certfile=None, ssl_context=None):
+                     certfile=None, ssl_context=None, timeout=None):
             if ssl_context is not None and keyfile is not None:
                 raise ValueError("ssl_context and keyfile arguments are mutually "
                                  "exclusive")
@@ -1294,20 +1304,20 @@ if HAVE_SSL:
                 ssl_context = ssl._create_stdlib_context(certfile=certfile,
                                                          keyfile=keyfile)
             self.ssl_context = ssl_context
-            IMAP4.__init__(self, host, port)
+            IMAP4.__init__(self, host, port, timeout)
 
-        def _create_socket(self):
-            sock = IMAP4._create_socket(self)
+        def _create_socket(self, timeout):
+            sock = IMAP4._create_socket(self, timeout)
             return self.ssl_context.wrap_socket(sock,
                                                 server_hostname=self.host)
 
-        def open(self, host='', port=IMAP4_SSL_PORT):
+        def open(self, host='', port=IMAP4_SSL_PORT, timeout=None):
             """Setup connection to remote server on "host:port".
                 (default: localhost:standard IMAP4 SSL port).
             This connection will be used by the routines:
                 read, readline, send, shutdown.
             """
-            IMAP4.open(self, host, port)
+            IMAP4.open(self, host, port, timeout)
 
     __all__.append("IMAP4_SSL")
 
@@ -1329,7 +1339,7 @@ class IMAP4_stream(IMAP4):
         IMAP4.__init__(self)
 
 
-    def open(self, host = None, port = None):
+    def open(self, host=None, port=None, timeout=None):
         """Setup a stream connection.
         This connection will be used by the routines:
             read, readline, send, shutdown.
diff --git a/Lib/test/test_imaplib.py b/Lib/test/test_imaplib.py
index 795276e..91aa771 100644
--- a/Lib/test/test_imaplib.py
+++ b/Lib/test/test_imaplib.py
@@ -440,6 +440,29 @@ class NewIMAPTestsMixin():
         with self.imap_class(*server.server_address):
             pass
 
+    def test_imaplib_timeout_test(self):
+        _, server = self._setup(SimpleIMAPHandler)
+        addr = server.server_address[1]
+        client = self.imap_class("localhost", addr, timeout=None)
+        self.assertEqual(client.sock.timeout, None)
+        client.shutdown()
+        client = self.imap_class("localhost", addr, timeout=support.LOOPBACK_TIMEOUT)
+        self.assertEqual(client.sock.timeout, support.LOOPBACK_TIMEOUT)
+        client.shutdown()
+        with self.assertRaises(ValueError):
+            client = self.imap_class("localhost", addr, timeout=0)
+
+    def test_imaplib_timeout_functionality_test(self):
+        class TimeoutHandler(SimpleIMAPHandler):
+            def handle(self):
+                time.sleep(1)
+                SimpleIMAPHandler.handle(self)
+
+        _, server = self._setup(TimeoutHandler)
+        addr = server.server_address[1]
+        with self.assertRaises(socket.timeout):
+            client = self.imap_class("localhost", addr, timeout=0.001)
+
     def test_with_statement(self):
         _, server = self._setup(SimpleIMAPHandler, connect=False)
         with self.imap_class(*server.server_address) as imap:
diff --git a/Misc/NEWS.d/next/Library/2019-11-17-17-32-35.bpo-38615.OVyaNX.rst b/Misc/NEWS.d/next/Library/2019-11-17-17-32-35.bpo-38615.OVyaNX.rst
new file mode 100644
index 0000000..04f51da
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2019-11-17-17-32-35.bpo-38615.OVyaNX.rst
@@ -0,0 +1,5 @@
+:class:`~imaplib.IMAP4` and :class:`~imaplib.IMAP4_SSL` now have an 
+optional *timeout* parameter for their constructors. 
+Also, the :meth:`~imaplib.IMAP4.open` method now has an optional *timeout* parameter
+with this change. The overridden methods of :class:`~imaplib.IMAP4_SSL` and
+:class:`~imaplib.IMAP4_stream` were applied to this change. Patch by Dong-hee Na.