Merge pull request #254 from shazow/factorout-connectionclsfactorout-connectioncls

Factor out HTTP(S)Connection -> ConnectionCls, and cleanup.

4 files changed, 150 insertions, 156 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index 93a7ceb8..042ea09b 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -7,6 +7,11 @@ dev (master)
 * Improved url parsing in ``urllib3.util.parse_url`` (properly parse '@' in
   username, and blank ports like 'hostname:').
 
+* New `urllib3.connection` module which contains all the HTTPConnection
+  objects.
+
+* ...
+
 
 1.7.1 (2013-09-25)
 ++++++++++++++++++
diff --git a/test/with_dummyserver/test_https.py b/test/with_dummyserver/test_https.py
index 8e07f2f1..199bc973 100644
--- a/test/with_dummyserver/test_https.py
+++ b/test/with_dummyserver/test_https.py
@@ -72,20 +72,14 @@ class TestHTTPS(HTTPSDummyServerTestCase):
             self.assertTrue("doesn't match" in str(e))
 
     def test_no_ssl(self):
-        import urllib3.connectionpool
-        OriginalHTTPSConnection = urllib3.connectionpool.HTTPSConnection
-        OriginalSSL = urllib3.connectionpool.ssl
-
-        urllib3.connectionpool.HTTPSConnection = None
-        urllib3.connectionpool.ssl = None
+        OriginalConnectionCls = self._pool.ConnectionCls
+        self._pool.ConnectionCls = None
 
         self.assertRaises(SSLError, self._pool._new_conn)
-
         self.assertRaises(SSLError, self._pool.request, 'GET', '/')
 
         # Undo
-        urllib3.HTTPSConnection = OriginalHTTPSConnection
-        urllib3.connectionpool.ssl = OriginalSSL
+        self._pool.ConnectionCls = OriginalConnectionCls
 
     def test_cert_reqs_as_constant(self):
         https_pool = HTTPSConnectionPool(self.host, self.port,
@@ -223,57 +217,28 @@ class TestHTTPS(HTTPSDummyServerTestCase):
 
 
     def test_enhanced_timeout(self):
-        import urllib3.connectionpool
-        OriginalHTTPSConnection = urllib3.connectionpool.HTTPSConnection
-        OriginalSSL = urllib3.connectionpool.ssl
+        def new_pool(timeout, cert_reqs='CERT_REQUIRED'):
+            https_pool = HTTPSConnectionPool(TARPIT_HOST, self.port,
+                                             timeout=timeout,
+                                             cert_reqs=cert_reqs)
+            return https_pool
 
-        urllib3.connectionpool.ssl = None
-
-        timeout = Timeout(connect=0.001)
-        https_pool = HTTPSConnectionPool(TARPIT_HOST, self.port,
-                                         timeout=timeout,
-                                         cert_reqs='CERT_REQUIRED')
+        https_pool = new_pool(Timeout(connect=0.001))
         conn = https_pool._new_conn()
-        self.assertEqual(conn.__class__, HTTPSConnection)
         self.assertRaises(ConnectTimeoutError, https_pool.request, 'GET', '/')
         self.assertRaises(ConnectTimeoutError, https_pool._make_request, conn,
                           'GET', '/')
 
-        timeout = Timeout(connect=5)
-        https_pool = HTTPSConnectionPool(TARPIT_HOST, self.port,
-                                         timeout=timeout,
-                                         cert_reqs='CERT_REQUIRED')
+        https_pool = new_pool(Timeout(connect=5))
         self.assertRaises(ConnectTimeoutError, https_pool.request, 'GET', '/',
                           timeout=Timeout(connect=0.001))
 
-        timeout = Timeout(total=None)
-        https_pool = HTTPSConnectionPool(TARPIT_HOST, self.port,
-                                         timeout=timeout,
-                                         cert_reqs='CERT_REQUIRED')
+        t = Timeout(total=None)
+        https_pool = new_pool(t)
         conn = https_pool._new_conn()
         self.assertRaises(ConnectTimeoutError, https_pool.request, 'GET', '/',
                           timeout=Timeout(total=None, connect=0.001))
 
-        https_pool = HTTPSConnectionPool(self.host, self.port,
-                                         timeout=timeout,
-                                         cert_reqs='CERT_NONE')
-        conn = https_pool._new_conn()
-        try:
-            conn.set_tunnel(self.host, self.port)
-        except AttributeError: # python 2.6
-            conn._set_tunnel(self.host, self.port)
-        conn._tunnel = mock.Mock()
-        try:
-            https_pool._make_request(conn, 'GET', '/')
-        except AttributeError:
-            # wrap_socket unavailable when you mock out ssl
-            pass
-        conn._tunnel.assert_called_once_with()
-
-        # Undo
-        urllib3.HTTPSConnection = OriginalHTTPSConnection
-        urllib3.connectionpool.ssl = OriginalSSL
-
     def test_enhanced_ssl_connection(self):
         conn = VerifiedHTTPSConnection(self.host, self.port)
         https_pool = HTTPSConnectionPool(self.host, self.port,
diff --git a/urllib3/connection.py b/urllib3/connection.py
new file mode 100644
index 00000000..e240786a
--- /dev/null
+++ b/urllib3/connection.py
@@ -0,0 +1,107 @@
+# urllib3/connection.py
+# Copyright 2008-2013 Andrey Petrov and contributors (see CONTRIBUTORS.txt)
+#
+# This module is part of urllib3 and is released under
+# the MIT License: http://www.opensource.org/licenses/mit-license.php
+
+import socket
+from socket import timeout as SocketTimeout
+
+try: # Python 3
+    from http.client import HTTPConnection, HTTPException
+except ImportError:
+    from httplib import HTTPConnection, HTTPException
+
+class DummyConnection(object):
+    "Used to detect a failed ConnectionCls import."
+    pass
+
+try: # Compiled with SSL?
+    ssl = None
+    HTTPSConnection = DummyConnection
+
+    class BaseSSLError(BaseException):
+        pass
+
+    try: # Python 3
+        from http.client import HTTPSConnection
+    except ImportError:
+        from httplib import HTTPSConnection
+
+    import ssl
+    BaseSSLError = ssl.SSLError
+
+except (ImportError, AttributeError): # Platform-specific: No SSL.
+    pass
+
+from .exceptions import (
+    ConnectTimeoutError,
+)
+from .packages.ssl_match_hostname import match_hostname
+from .util import (
+    assert_fingerprint,
+    resolve_cert_reqs,
+    resolve_ssl_version,
+    ssl_wrap_socket,
+)
+
+class VerifiedHTTPSConnection(HTTPSConnection):
+    """
+    Based on httplib.HTTPSConnection but wraps the socket with
+    SSL certification.
+    """
+    cert_reqs = None
+    ca_certs = None
+    ssl_version = None
+
+    def set_cert(self, key_file=None, cert_file=None,
+                 cert_reqs=None, ca_certs=None,
+                 assert_hostname=None, assert_fingerprint=None):
+
+        self.key_file = key_file
+        self.cert_file = cert_file
+        self.cert_reqs = cert_reqs
+        self.ca_certs = ca_certs
+        self.assert_hostname = assert_hostname
+        self.assert_fingerprint = assert_fingerprint
+
+    def connect(self):
+        # Add certificate verification
+        try:
+            sock = socket.create_connection(
+                address=(self.host, self.port),
+                timeout=self.timeout,
+            )
+        except SocketTimeout:
+                raise ConnectTimeoutError(
+                    self, "Connection to %s timed out. (connect timeout=%s)" %
+                    (self.host, self.timeout))
+
+        resolved_cert_reqs = resolve_cert_reqs(self.cert_reqs)
+        resolved_ssl_version = resolve_ssl_version(self.ssl_version)
+
+        if self._tunnel_host:
+            self.sock = sock
+            # Calls self._set_hostport(), so self.host is
+            # self._tunnel_host below.
+            self._tunnel()
+
+        # Wrap socket using verification with the root certs in
+        # trusted_root_certs
+        self.sock = ssl_wrap_socket(sock, self.key_file, self.cert_file,
+                                    cert_reqs=resolved_cert_reqs,
+                                    ca_certs=self.ca_certs,
+                                    server_hostname=self.host,
+                                    ssl_version=resolved_ssl_version)
+
+        if resolved_cert_reqs != ssl.CERT_NONE:
+            if self.assert_fingerprint:
+                assert_fingerprint(self.sock.getpeercert(binary_form=True),
+                                   self.assert_fingerprint)
+            elif self.assert_hostname is not False:
+                match_hostname(self.sock.getpeercert(),
+                               self.assert_hostname or self.host)
+
+
+if ssl:
+    HTTPSConnection = VerifiedHTTPSConnection
diff --git a/urllib3/connectionpool.py b/urllib3/connectionpool.py
index 1ad20a24..cbd89eed 100644
--- a/urllib3/connectionpool.py
+++ b/urllib3/connectionpool.py
@@ -11,39 +11,12 @@ from socket import error as SocketError, timeout as SocketTimeout
 import socket
 
 try: # Python 3
-    from http.client import HTTPConnection, HTTPException
-    from http.client import HTTP_PORT, HTTPS_PORT
-except ImportError:
-    from httplib import HTTPConnection, HTTPException
-    from httplib import HTTP_PORT, HTTPS_PORT
-
-try: # Python 3
     from queue import LifoQueue, Empty, Full
 except ImportError:
     from Queue import LifoQueue, Empty, Full
     import Queue as _  # Platform-specific: Windows
 
 
-try: # Compiled with SSL?
-    HTTPSConnection = object
-
-    class BaseSSLError(Exception):
-        pass
-
-    ssl = None
-
-    try: # Python 3
-        from http.client import HTTPSConnection
-    except ImportError:
-        from httplib import HTTPSConnection
-
-    import ssl
-    BaseSSLError = ssl.SSLError
-
-except (ImportError, AttributeError): # Platform-specific: No SSL.
-    pass
-
-
 from .exceptions import (
     ClosedPoolError,
     ConnectTimeoutError,
@@ -54,20 +27,23 @@ from .exceptions import (
     ReadTimeoutError,
     ProxyError,
 )
-from .packages.ssl_match_hostname import CertificateError, match_hostname
+from .packages.ssl_match_hostname import CertificateError
 from .packages import six
+from .connection import (
+    DummyConnection,
+    HTTPConnection, HTTPSConnection, VerifiedHTTPSConnection,
+    HTTPException, BaseSSLError,
+)
 from .request import RequestMethods
 from .response import HTTPResponse
 from .util import (
     assert_fingerprint,
     get_host,
     is_connection_dropped,
-    resolve_cert_reqs,
-    resolve_ssl_version,
-    ssl_wrap_socket,
     Timeout,
 )
 
+
 xrange = six.moves.xrange
 
 log = logging.getLogger(__name__)
@@ -75,70 +51,11 @@ log = logging.getLogger(__name__)
 _Default = object()
 
 port_by_scheme = {
-    'http': HTTP_PORT,
-    'https': HTTPS_PORT,
+    'http': 80,
+    'https': 443,
 }
 
 
-## Connection objects (extension of httplib)
-
-class VerifiedHTTPSConnection(HTTPSConnection):
-    """
-    Based on httplib.HTTPSConnection but wraps the socket with
-    SSL certification.
-    """
-    cert_reqs = None
-    ca_certs = None
-    ssl_version = None
-
-    def set_cert(self, key_file=None, cert_file=None,
-                 cert_reqs=None, ca_certs=None,
-                 assert_hostname=None, assert_fingerprint=None):
-
-        self.key_file = key_file
-        self.cert_file = cert_file
-        self.cert_reqs = cert_reqs
-        self.ca_certs = ca_certs
-        self.assert_hostname = assert_hostname
-        self.assert_fingerprint = assert_fingerprint
-
-    def connect(self):
-        # Add certificate verification
-        try:
-            sock = socket.create_connection(
-                address=(self.host, self.port),
-                timeout=self.timeout)
-        except SocketTimeout:
-                raise ConnectTimeoutError(
-                    self, "Connection to %s timed out. (connect timeout=%s)" %
-                    (self.host, self.timeout))
-
-        resolved_cert_reqs = resolve_cert_reqs(self.cert_reqs)
-        resolved_ssl_version = resolve_ssl_version(self.ssl_version)
-
-        if self._tunnel_host:
-            self.sock = sock
-            # Calls self._set_hostport(), so self.host is
-            # self._tunnel_host below.
-            self._tunnel()
-
-        # Wrap socket using verification with the root certs in
-        # trusted_root_certs
-        self.sock = ssl_wrap_socket(sock, self.key_file, self.cert_file,
-                                    cert_reqs=resolved_cert_reqs,
-                                    ca_certs=self.ca_certs,
-                                    server_hostname=self.host,
-                                    ssl_version=resolved_ssl_version)
-
-        if resolved_cert_reqs != ssl.CERT_NONE:
-            if self.assert_fingerprint:
-                assert_fingerprint(self.sock.getpeercert(binary_form=True),
-                                   self.assert_fingerprint)
-            elif self.assert_hostname is not False:
-                match_hostname(self.sock.getpeercert(),
-                               self.assert_hostname or self.host)
-
-
 ## Pool objects
 
 class ConnectionPool(object):
@@ -218,6 +135,7 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods):
     """
 
     scheme = 'http'
+    ConnectionCls = HTTPConnection
 
     def __init__(self, host, port=None, strict=False,
                  timeout=Timeout.DEFAULT_TIMEOUT, maxsize=1, block=False,
@@ -255,14 +173,14 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods):
         self.num_connections += 1
         log.info("Starting new HTTP connection (%d): %s" %
                  (self.num_connections, self.host))
+
         extra_params = {}
         if not six.PY3:  # Python 2
             extra_params['strict'] = self.strict
 
-        return HTTPConnection(host=self.host, port=self.port,
-                              timeout=self.timeout.connect_timeout,
-                              **extra_params)
-
+        return self.ConnectionCls(host=self.host, port=self.port,
+                                  timeout=self.timeout.connect_timeout,
+                                  **extra_params)
 
     def _get_conn(self, timeout=None):
         """
@@ -362,7 +280,7 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods):
             timeout_obj.start_connect()
             conn.timeout = timeout_obj.connect_timeout
             # conn.request() calls httplib.*.request, not the method in
-            # request.py. It also calls makefile (recv) on the socket
+            # urllib3.request. It also calls makefile (recv) on the socket.
             conn.request(method, url, **httplib_request_kw)
         except SocketTimeout:
             raise ConnectTimeoutError(
@@ -371,7 +289,7 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods):
 
         # Reset the timeout for the recv() on the socket
         read_timeout = timeout_obj.read_timeout
-        log.debug("Setting read timeout to %s" % read_timeout)
+
         # App Engine doesn't have a sock attr
         if hasattr(conn, 'sock') and \
             read_timeout is not None and \
@@ -639,6 +557,7 @@ class HTTPSConnectionPool(HTTPConnectionPool):
     """
 
     scheme = 'https'
+    ConnectionCls = HTTPSConnection
 
     def __init__(self, host, port=None,
                  strict=False, timeout=None, maxsize=1,
@@ -694,26 +613,24 @@ class HTTPSConnectionPool(HTTPConnectionPool):
         log.info("Starting new HTTPS connection (%d): %s"
                  % (self.num_connections, self.host))
 
+        if not self.ConnectionCls or self.ConnectionCls is DummyConnection:
+            # Platform-specific: Python without ssl
+            raise SSLError("Can't connect to HTTPS URL because the SSL "
+                           "module is not available.")
+
         actual_host = self.host
         actual_port = self.port
         if self.proxy is not None:
             actual_host = self.proxy.host
             actual_port = self.proxy.port
 
-        if not ssl:  # Platform-specific: Python compiled without +ssl
-            if not HTTPSConnection or HTTPSConnection is object:
-                raise SSLError("Can't connect to HTTPS URL because the SSL "
-                               "module is not available.")
-            connection_class = HTTPSConnection
-        else:
-            connection_class = VerifiedHTTPSConnection
-
         extra_params = {}
         if not six.PY3:  # Python 2
             extra_params['strict'] = self.strict
-        conn = connection_class(host=actual_host, port=actual_port,
-                                timeout=self.timeout.connect_timeout,
-                                **extra_params)
+
+        conn = self.ConnectionCls(host=actual_host, port=actual_port,
+                                  timeout=self.timeout.connect_timeout,
+                                  **extra_params)
 
         return self._prepare_conn(conn)