prevent HEAD requests from auto-closing the connection

httplib2_noclose_HEAD.py

import httplib
import socket

import httplib2

# the whole reason this module is necessary is to prevent HEAD requests from auto-closing the connection.

class HTTPResponse(httplib.HTTPResponse):
    def read(self, amt=None):
        # prevent the HEAD check from closing the connection
        if self._method == 'HEAD':
            self._method = 'GET'
            rsp = httplib.HTTPResponse.read(self, amt)
            self._method = 'HEAD'
            return rsp
        else:
            return httplib.HTTPResponse.read(self, amt)

class do_not_close_http(httplib2.HTTPConnectionWithTimeout):
    response_class = HTTPResponse

class do_not_close_https(httplib2.HTTPSConnectionWithTimeout):
    response_class = HTTPResponse
class Http(httplib2.Http):

    def __init__(self, cache=None, timeout=None, proxy_info=httplib2.proxy_info_from_environment, ca_certs=None,
                 disable_ssl_certificate_validation=False, ssl_version=None):
        super(Http, self).__init__(cache, timeout, proxy_info, ca_certs, disable_ssl_certificate_validation,
                                   ssl_version)
        self.force_exception_to_status_code = True

    def request(self, uri, method="GET", body=None, headers=None, redirections=httplib2.DEFAULT_MAX_REDIRECTS,connection_type=None):

        return super(Http, self).request(uri, method, body, headers, redirections, connection_type=connection_type or self.get_connection_type(uri))

    def get_connection_type(self,uri):
        scheme, authority, path, query, fragment = httplib2.parse_uri(uri)
        return do_not_close_http if scheme.lower() == 'http' else do_not_close_https

    def _conn_request(self, conn, request_uri, method, body, headers):
        # do not close on HEAD!
        RETRIES=httplib2.RETRIES
        ServerNotFoundError=httplib2.ServerNotFoundError
        errno=httplib2.errno
        ssl_SSLError=httplib2.ssl_SSLError
        # noinspection PyProtectedMember
        _decompressContent=httplib2._decompressContent
        Response=httplib2.Response
        response=None
        content=None
        # log.info('using conn: %s',conn.sock and conn.sock.getsockname())
        i = 0
        seen_bad_status_line = False
        while i < RETRIES:
            i += 1
            try:
                if hasattr(conn, 'sock') and conn.sock is None:
                    conn.connect()
                conn.request(method, request_uri, body, headers)
            except socket.timeout:
                raise
            except socket.gaierror:
                conn.close()
                raise ServerNotFoundError("Unable to find the server at %s" % conn.host)
            except ssl_SSLError:
                conn.close()
                raise
            except socket.error, e:
                if hasattr(e, 'args'):
                    err = getattr(e, 'args')[0]
                else:
                    err = e.errno
                if err == errno.ECONNREFUSED: # Connection refused
                    raise
                if err in (errno.ENETUNREACH, errno.EADDRNOTAVAIL) and i < RETRIES:
                    continue  # retry on potentially transient socket errors
            except httplib.HTTPException:
                # Just because the server closed the connection doesn't apparently mean
                # that the server didn't send a response.
                if hasattr(conn, 'sock') and conn.sock is None:
                    if i < RETRIES-1:
                        conn.close()
                        conn.connect()
                        continue
                    else:
                        conn.close()
                        raise
                if i < RETRIES-1:
                    conn.close()
                    conn.connect()
                    continue
            try:
                response = conn.getresponse()
            except httplib.BadStatusLine:
                # If we get a BadStatusLine on the first try then that means
                # the connection just went stale, so retry regardless of the
                # number of RETRIES set.
                if not seen_bad_status_line and i == 1:
                    i = 0
                    seen_bad_status_line = True
                    conn.close()
                    conn.connect()
                    continue
                else:
                    conn.close()
                    raise
            except (socket.error, httplib.HTTPException):
                if i < RETRIES-1:
                    conn.close()
                    conn.connect()
                    continue
                else:
                    conn.close()
                    raise
            else:
                content = response.read()
                response = Response(response)
                if method != "HEAD":
                    content = _decompressContent(response, content)
            break
        return response, content