Request to allow passing CustomBackend to httpx client

[Cosmicoppai]

Is there a way to pass custom backend to httpx Client. After looking at the code, it seems there is no way until I override the whole chain of class.

It would be useful to have something like:

client = Client(network_backend)
# the network_backend is of type httpcore._backend.base.NetworkBackend

# An example:
from httpx._client import Client, EventHook
from httpx._config import DEFAULT_LIMITS, Limits, Proxy, DEFAULT_TIMEOUT_CONFIG, DEFAULT_MAX_REDIRECTS
from httpx._transports.base import BaseTransport
from httpx._transports.wsgi import WSGITransport
from httpcore import NetworkBackend
from httpx._types import (
    AuthTypes,
    CertTypes,
    CookieTypes,
    HeaderTypes,
    ProxiesTypes,
    ProxyTypes,
    QueryParamTypes,
    TimeoutTypes,
    URLTypes,
    VerifyTypes
)

from .transport import CustomHTTPTransport
import typing
from .utils import DataCounter


class CustClient(Client):
    def __init__(
            self,
            *,
            auth: AuthTypes | None = None,
            params: QueryParamTypes | None = None,
            headers: HeaderTypes | None = None,
            cookies: CookieTypes | None = None,
            verify: VerifyTypes = True,
            cert: CertTypes | None = None,
            http1: bool = True,
            http2: bool = False,
            proxy: ProxyTypes | None = None,
            proxies: ProxiesTypes | None = None,
            mounts: None | (typing.Mapping[str, BaseTransport | None]) = None,
            timeout: TimeoutTypes = DEFAULT_TIMEOUT_CONFIG,
            follow_redirects: bool = False,
            limits: Limits = DEFAULT_LIMITS,
            max_redirects: int = DEFAULT_MAX_REDIRECTS,
            event_hooks: None | (typing.Mapping[str, list[EventHook]]) = None,
            base_url: URLTypes = "",
            transport: BaseTransport | None = None,
            app: typing.Callable[..., typing.Any] | None = None,
            trust_env: bool = True,
            default_encoding: str | typing.Callable[[bytes], str] = "utf-8",
            network_backend: NetworkBackend | None = None
    ) -> None:
        self._network_backend = network_backend
        self._data_counter: DataCounter = DataCounter()
        super().__init__(
            auth=auth, params=params, headers=headers, cookies=cookies, verify=verify,
            cert=cert, http1=http1, http2=http2, proxy=proxy, proxies=proxies,
            mounts=mounts, timeout=timeout, follow_redirects=follow_redirects,
            limits=limits, max_redirects=max_redirects, event_hooks=event_hooks,
            base_url=base_url, transport=transport, app=app, trust_env=trust_env,
            default_encoding=default_encoding
        )

    @property
    def sent(self):
        return self._data_counter.sent

    @property
    def received(self):
        return self._data_counter.received

    def _init_transport(
            self,
            verify: VerifyTypes = True,
            cert: CertTypes | None = None,
            http1: bool = True,
            http2: bool = False,
            limits: Limits = DEFAULT_LIMITS,
            transport: BaseTransport | None = None,
            app: typing.Callable[..., typing.Any] | None = None,
            trust_env: bool = True,
    ) -> BaseTransport:
        if transport is not None:
            return transport

        if app is not None:
            return WSGITransport(app=app)

        return CustomHTTPTransport(
            verify=verify,
            cert=cert,
            http1=http1,
            http2=http2,
            limits=limits,
            trust_env=trust_env,
            network_backend=self._network_backend,
            counter=self._data_counter
        )

    def _init_proxy_transport(
            self,
            proxy: Proxy,
            verify: VerifyTypes = True,
            cert: CertTypes | None = None,
            http1: bool = True,
            http2: bool = False,
            limits: Limits = DEFAULT_LIMITS,
            trust_env: bool = True,
    ) -> BaseTransport:
        return CustomHTTPTransport(
            verify=verify,
            cert=cert,
            http1=http1,
            http2=http2,
            limits=limits,
            trust_env=trust_env,
            proxy=proxy,
            network_backend=self._network_backend,
            counter=self._data_counter
        )

import httpcore
import typing

from httpx._transports.default import HTTPTransport, ResponseStream, map_httpcore_exceptions
from httpx._config import DEFAULT_LIMITS, Limits, Proxy, create_ssl_context
from httpx._types import CertTypes, ProxyTypes, VerifyTypes, SyncByteStream
from httpx._urls import URL
from httpx._models import Request, Response
from .backend import SyncCustBackend
from .utils import DataCounter

T = typing.TypeVar("T", bound="HTTPTransport")
A = typing.TypeVar("A", bound="AsyncHTTPTransport")

SOCKET_OPTION = typing.Union[
    typing.Tuple[int, int, int],
    typing.Tuple[int, int, typing.Union[bytes, bytearray]],
    typing.Tuple[int, int, None, int],
]


class CustomHTTPTransport(HTTPTransport):
    def __init__(
            self,
            verify: VerifyTypes = True,
            cert: CertTypes | None = None,
            http1: bool = True,
            http2: bool = False,
            limits: Limits = DEFAULT_LIMITS,
            trust_env: bool = True,
            proxy: ProxyTypes | None = None,
            uds: str | None = None,
            local_address: str | None = None,
            retries: int = 0,
            socket_options: typing.Iterable[SOCKET_OPTION] | None = None,
            network_backend: httpcore.NetworkBackend | None = None,
            counter: DataCounter | None = None
    ) -> None:
        super().__init__(verify, cert, http1, http2, limits, trust_env, proxy, uds,
                         local_address, retries, socket_options)

        self._counter = counter

        if network_backend is None:
            network_backend = SyncCustBackend(self._counter)

        ssl_context = create_ssl_context(verify=verify, cert=cert, trust_env=trust_env)
        proxy = Proxy(url=proxy) if isinstance(proxy, (str, URL)) else proxy

        if proxy is None:
            self._pool = httpcore.ConnectionPool(
                ssl_context=ssl_context,
                max_connections=limits.max_connections,
                max_keepalive_connections=limits.max_keepalive_connections,
                keepalive_expiry=limits.keepalive_expiry,
                http1=http1,
                http2=http2,
                uds=uds,
                local_address=local_address,
                retries=retries,
                socket_options=socket_options,
                network_backend=network_backend
            )
        elif proxy.url.scheme in ("http", "https"):
            self._pool = httpcore.HTTPProxy(
                proxy_url=httpcore.URL(
                    scheme=proxy.url.raw_scheme,
                    host=proxy.url.raw_host,
                    port=proxy.url.port,
                    target=proxy.url.raw_path
                ),
                proxy_auth=proxy.raw_auth,
                proxy_headers=proxy.headers.raw,
                ssl_context=ssl_context,
                proxy_ssl_context=proxy.ssl_context,
                max_connections=limits.max_connections,
                max_keepalive_connections=limits.max_keepalive_connections,
                keepalive_expiry=limits.keepalive_expiry,
                http1=http1,
                http2=http2,
                socket_options=socket_options,
                network_backend=network_backend
            )
        elif proxy.url.scheme == "socks5":
            try:
                import socksio  # noqa
            except ImportError:  # pragma: no cover
                raise ImportError(
                    "Using SOCKS proxy, but the 'socksio' package is not installed. "
                    "Make sure to install httpx using `pip install httpx[socks]`."
                ) from None

            self._pool = httpcore.SOCKSProxy(
                proxy_url=httpcore.URL(
                    scheme=proxy.url.raw_scheme,
                    host=proxy.url.raw_host,
                    port=proxy.url.port,
                    target=proxy.url.raw_path,
                ),
                proxy_auth=proxy.raw_auth,
                ssl_context=ssl_context,
                max_connections=limits.max_connections,
                max_keepalive_connections=limits.max_keepalive_connections,
                keepalive_expiry=limits.keepalive_expiry,
                http1=http1,
                http2=http2,
                network_backend=network_backend
            )
        else:  # pragma: no cover
            raise ValueError(
                "Proxy protocol must be either 'http', 'https', or 'socks5',"
                f" but got {proxy.url.scheme!r}."
            )

    def handle_request(self, request: Request) -> Response:
        assert isinstance(request.stream, SyncByteStream)

        req = httpcore.Request(
            method=request.method,
            url=httpcore.URL(
                scheme=request.url.raw_scheme,
                host=request.url.raw_host,
                port=request.url.port,
                target=request.url.raw_path,
            ),
            headers=request.headers.raw,
            content=request.stream,
            extensions=request.extensions,
        )

        # reset counter before sending the data
        # not every request will contain tls negotiation data, as there is connection pool
        self._reset_counter()
        with map_httpcore_exceptions():
            resp = self._pool.handle_request(req)

        assert isinstance(resp.stream, typing.Iterable)

        return Response(
            status_code=resp.status,
            headers=resp.headers,
            stream=ResponseStream(resp.stream),
            extensions=resp.extensions,
        )

    def _reset_counter(self) -> None:
        self._counter.sent = 0
        self._counter.received = 0

import socket
import ssl
import sys
import typing
from httpcore._exceptions import (
    ConnectError,
    ConnectTimeout,
    ExceptionMapping,
    map_exceptions,
)
from httpcore._backends.base import SOCKET_OPTION, NetworkBackend, NetworkStream
from httpcore._backends.sync import TLSinTLSStream, SyncStream, SyncBackend
from .utils import DataCounter


class LoggingSSLWrapper:
    def __init__(self, wrapped_socket: ssl.SSLSocket, counter: DataCounter):
        self._wrapped_socket = wrapped_socket
        self._counter = counter

    def send(self, data, flags=0):
        self._counter.sent += len(data)
        return self._wrapped_socket.send(data, flags)

    def recv(self, bufsize, flags=0):
        data = self._wrapped_socket.recv(bufsize, flags)
        self._counter.received += len(data)
        return data

    def sendall(self, data, flags=0):
        self._counter.sent = len(data)
        return self._wrapped_socket.sendall(data, flags)

    def close(self):
        return self._wrapped_socket.close()

    def do_handshake(self):
        self._wrapped_socket.do_handshake()

    def __getattr__(self, attr):
        return getattr(self._wrapped_socket, attr)


class SyncCustStream(SyncStream):
    def __init__(self, sock: socket.socket, counter: DataCounter) -> None:
        super().__init__(sock)
        self._counter = counter

    def start_tls(
        self,
        ssl_context: ssl.SSLContext,
        server_hostname: typing.Optional[str] = None,
        timeout: typing.Optional[float] = None,
    ) -> NetworkStream:
        exc_map: ExceptionMapping = {
            socket.timeout: ConnectTimeout,
            OSError: ConnectError,
        }
        with map_exceptions(exc_map):
            try:
                if isinstance(self._sock, ssl.SSLSocket):  # pragma: no cover
                    # If the underlying socket has already been upgraded
                    # to the TLS layer (i.e. is an instance of SSLSocket),
                    # we need some additional smarts to support TLS-in-TLS.
                    return TLSinTLSStream(
                        self._sock, ssl_context, server_hostname, timeout
                    )
                else:
                    self._sock.settimeout(timeout)
                    ssl_sock = ssl_context.wrap_socket(
                        self._sock, server_hostname=server_hostname
                    )
                    logging_ssl_sock = LoggingSSLWrapper(ssl_sock, self._counter)
                    # logging_ssl_sock.do_handshake()
            except Exception as exc:  # pragma: nocover
                self.close()
                raise exc
        return SyncCustStream(logging_ssl_sock, self._counter)


class SyncCustBackend(SyncBackend):
    def __init__(self, counter: DataCounter):
        self._counter = counter

    def connect_tcp(
        self,
        host: str,
        port: int,
        timeout: typing.Optional[float] = None,
        local_address: typing.Optional[str] = None,
        socket_options: typing.Optional[typing.Iterable[SOCKET_OPTION]] = None,
    ) -> NetworkStream:
        # Note that we automatically include `TCP_NODELAY`
        # in addition to any other custom socket options.
        if socket_options is None:
            socket_options = []  # pragma: no cover
        address = (host, port)
        source_address = None if local_address is None else (local_address, 0)
        exc_map: ExceptionMapping = {
            socket.timeout: ConnectTimeout,
            OSError: ConnectError,
        }

        with map_exceptions(exc_map):
            sock = socket.create_connection(
                address,
                timeout,
                source_address=source_address,
            )
            for option in socket_options:
                sock.setsockopt(*option)  # pragma: no cover
            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        return SyncCustStream(sock, self._data_counter)

    def connect_unix_socket(
        self,
        path: str,
        timeout: typing.Optional[float] = None,
        socket_options: typing.Optional[typing.Iterable[SOCKET_OPTION]] = None,
    ) -> NetworkStream:  # pragma: nocover
        if sys.platform == "win32":
            raise RuntimeError(
                "Attempted to connect to a UNIX socket on a Windows system."
            )
        if socket_options is None:
            socket_options = []

        exc_map: ExceptionMapping = {
            socket.timeout: ConnectTimeout,
            OSError: ConnectError,
        }
        with map_exceptions(exc_map):
            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            for option in socket_options:
                sock.setsockopt(*option)
            sock.settimeout(timeout)
            sock.connect(path)
        return SyncCustStream(sock, self._data_counter)

[tomchristie]

Yep, supporting custom backends to httpx.HTTPTransport and httpx.AsyncHTTPTransport would be a good call.

• httpcore supports custom backends being passed to the connection pool.
• httpx does not currently have any way of passing them through.

Perhaps you could start with a docs-first PR, and add in a network_backend = ... example, similar to the existing uds docs.