Add support for socket options

httpcore/_async/connection.py

@@ -1,15 +1,15 @@
 import itertools
 import ssl
 from types import TracebackType
-from typing import Iterator, Optional, Type
+from typing import Iterable, Iterator, Optional, Type
 
 from .._exceptions import ConnectError, ConnectionNotAvailable, ConnectTimeout
 from .._models import Origin, Request, Response
 from .._ssl import default_ssl_context
 from .._synchronization import AsyncLock
 from .._trace import Trace
 from ..backends.auto import AutoBackend
-from ..backends.base import AsyncNetworkBackend, AsyncNetworkStream
+from ..backends.base import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
 from .http11 import AsyncHTTP11Connection
 from .interfaces import AsyncConnectionInterface
 
@@ -34,6 +34,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._origin = origin
         self._ssl_context = ssl_context
@@ -50,6 +51,7 @@ def __init__(
         self._connection: Optional[AsyncConnectionInterface] = None
         self._connect_failed: bool = False
         self._request_lock = AsyncLock()
+        self._socket_options = () if socket_options is None else socket_options
 
     async def handle_async_request(self, request: Request) -> Response:
         if not self.can_handle_request(request.url.origin):
@@ -104,6 +106,7 @@ async def _connect(self, request: Request) -> AsyncNetworkStream:
                         "port": self._origin.port,
                         "local_address": self._local_address,
                         "timeout": timeout,
+                        "socket_options": self._socket_options,
                     }
                     async with Trace(
                         "connection.connect_tcp", request, kwargs
@@ -114,6 +117,7 @@ async def _connect(self, request: Request) -> AsyncNetworkStream:
                     kwargs = {
                         "path": self._uds,
                         "timeout": timeout,
+                        "socket_options": self._socket_options,
                     }
                     async with Trace(
                         "connection.connect_unix_socket", request, kwargs

httpcore/_async/connection_pool.py

@@ -1,13 +1,13 @@
 import ssl
 import sys
 from types import TracebackType
-from typing import AsyncIterable, AsyncIterator, List, Optional, Type
+from typing import AsyncIterable, AsyncIterator, Iterable, List, Optional, Type
 
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
 from .._models import Origin, Request, Response
 from .._synchronization import AsyncEvent, AsyncLock
 from ..backends.auto import AutoBackend
-from ..backends.base import AsyncNetworkBackend
+from ..backends.base import SOCKET_OPTION, AsyncNetworkBackend
 from .connection import AsyncHTTPConnection
 from .interfaces import AsyncConnectionInterface, AsyncRequestInterface
 
@@ -53,6 +53,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         """
         A connection pool for making HTTP requests.
@@ -108,6 +109,7 @@ def __init__(
         self._network_backend = (
             AutoBackend() if network_backend is None else network_backend
         )
+        self._socket_options = () if socket_options is None else socket_options
 
     def create_connection(self, origin: Origin) -> AsyncConnectionInterface:
         return AsyncHTTPConnection(
@@ -120,6 +122,7 @@ def create_connection(self, origin: Origin) -> AsyncConnectionInterface:
             local_address=self._local_address,
             uds=self._uds,
             network_backend=self._network_backend,
+            socket_options=self._socket_options,
         )
 
     @property

httpcore/_async/http_proxy.py

@@ -1,6 +1,6 @@
 import ssl
 from base64 import b64encode
-from typing import List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Iterable, List, Mapping, Optional, Sequence, Tuple, Union
 
 from .._exceptions import ProxyError
 from .._models import (
@@ -15,7 +15,7 @@
 from .._ssl import default_ssl_context
 from .._synchronization import AsyncLock
 from .._trace import Trace
-from ..backends.base import AsyncNetworkBackend
+from ..backends.base import SOCKET_OPTION, AsyncNetworkBackend
 from .connection import AsyncHTTPConnection
 from .connection_pool import AsyncConnectionPool
 from .http11 import AsyncHTTP11Connection
@@ -69,6 +69,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         """
         A connection pool for making HTTP requests.
@@ -115,6 +116,7 @@ def __init__(
             retries=retries,
             local_address=local_address,
             uds=uds,
+            socket_options=socket_options,
         )
         self._ssl_context = ssl_context
         self._proxy_url = enforce_url(proxy_url, name="proxy_url")
@@ -156,11 +158,13 @@ def __init__(
         proxy_headers: Union[HeadersAsMapping, HeadersAsSequence, None] = None,
         keepalive_expiry: Optional[float] = None,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._connection = AsyncHTTPConnection(
             origin=proxy_origin,
             keepalive_expiry=keepalive_expiry,
             network_backend=network_backend,
+            socket_options=socket_options,
         )
         self._proxy_origin = proxy_origin
         self._proxy_headers = enforce_headers(proxy_headers, name="proxy_headers")
@@ -219,11 +223,13 @@ def __init__(
         http1: bool = True,
         http2: bool = False,
         network_backend: Optional[AsyncNetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._connection: AsyncConnectionInterface = AsyncHTTPConnection(
             origin=proxy_origin,
             keepalive_expiry=keepalive_expiry,
             network_backend=network_backend,
+            socket_options=socket_options,
         )
         self._proxy_origin = proxy_origin
         self._remote_origin = remote_origin

httpcore/_sync/connection.py

@@ -1,15 +1,15 @@
 import itertools
 import ssl
 from types import TracebackType
-from typing import Iterator, Optional, Type
+from typing import Iterable, Iterator, Optional, Type
 
 from .._exceptions import ConnectError, ConnectionNotAvailable, ConnectTimeout
 from .._models import Origin, Request, Response
 from .._ssl import default_ssl_context
 from .._synchronization import Lock
 from .._trace import Trace
 from ..backends.sync import SyncBackend
-from ..backends.base import NetworkBackend, NetworkStream
+from ..backends.base import SOCKET_OPTION, NetworkBackend, NetworkStream
 from .http11 import HTTP11Connection
 from .interfaces import ConnectionInterface
 
@@ -34,6 +34,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._origin = origin
         self._ssl_context = ssl_context
@@ -50,6 +51,7 @@ def __init__(
         self._connection: Optional[ConnectionInterface] = None
         self._connect_failed: bool = False
         self._request_lock = Lock()
+        self._socket_options = () if socket_options is None else socket_options
 
     def handle_request(self, request: Request) -> Response:
         if not self.can_handle_request(request.url.origin):
@@ -104,6 +106,7 @@ def _connect(self, request: Request) -> NetworkStream:
                         "port": self._origin.port,
                         "local_address": self._local_address,
                         "timeout": timeout,
+                        "socket_options": self._socket_options,
                     }
                     with Trace(
                         "connection.connect_tcp", request, kwargs
@@ -114,6 +117,7 @@ def _connect(self, request: Request) -> NetworkStream:
                     kwargs = {
                         "path": self._uds,
                         "timeout": timeout,
+                        "socket_options": self._socket_options,
                     }
                     with Trace(
                         "connection.connect_unix_socket", request, kwargs

httpcore/_sync/connection_pool.py

@@ -1,13 +1,13 @@
 import ssl
 import sys
 from types import TracebackType
-from typing import Iterable, Iterator, List, Optional, Type
+from typing import Iterable, Iterator, Iterable, List, Optional, Type
 
 from .._exceptions import ConnectionNotAvailable, UnsupportedProtocol
 from .._models import Origin, Request, Response
 from .._synchronization import Event, Lock
 from ..backends.sync import SyncBackend
-from ..backends.base import NetworkBackend
+from ..backends.base import SOCKET_OPTION, NetworkBackend
 from .connection import HTTPConnection
 from .interfaces import ConnectionInterface, RequestInterface
 
@@ -53,6 +53,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         """
         A connection pool for making HTTP requests.
@@ -108,6 +109,7 @@ def __init__(
         self._network_backend = (
             SyncBackend() if network_backend is None else network_backend
         )
+        self._socket_options = () if socket_options is None else socket_options
 
     def create_connection(self, origin: Origin) -> ConnectionInterface:
         return HTTPConnection(
@@ -120,6 +122,7 @@ def create_connection(self, origin: Origin) -> ConnectionInterface:
             local_address=self._local_address,
             uds=self._uds,
             network_backend=self._network_backend,
+            socket_options=self._socket_options,
         )
 
     @property

httpcore/_sync/http_proxy.py

@@ -1,6 +1,6 @@
 import ssl
 from base64 import b64encode
-from typing import List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Iterable, List, Mapping, Optional, Sequence, Tuple, Union
 
 from .._exceptions import ProxyError
 from .._models import (
@@ -15,7 +15,7 @@
 from .._ssl import default_ssl_context
 from .._synchronization import Lock
 from .._trace import Trace
-from ..backends.base import NetworkBackend
+from ..backends.base import SOCKET_OPTION, NetworkBackend
 from .connection import HTTPConnection
 from .connection_pool import ConnectionPool
 from .http11 import HTTP11Connection
@@ -69,6 +69,7 @@ def __init__(
         local_address: Optional[str] = None,
         uds: Optional[str] = None,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         """
         A connection pool for making HTTP requests.
@@ -115,6 +116,7 @@ def __init__(
             retries=retries,
             local_address=local_address,
             uds=uds,
+            socket_options=socket_options,
         )
         self._ssl_context = ssl_context
         self._proxy_url = enforce_url(proxy_url, name="proxy_url")
@@ -156,11 +158,13 @@ def __init__(
         proxy_headers: Union[HeadersAsMapping, HeadersAsSequence, None] = None,
         keepalive_expiry: Optional[float] = None,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._connection = HTTPConnection(
             origin=proxy_origin,
             keepalive_expiry=keepalive_expiry,
             network_backend=network_backend,
+            socket_options=socket_options,
         )
         self._proxy_origin = proxy_origin
         self._proxy_headers = enforce_headers(proxy_headers, name="proxy_headers")
@@ -219,11 +223,13 @@ def __init__(
         http1: bool = True,
         http2: bool = False,
         network_backend: Optional[NetworkBackend] = None,
+        socket_options: Optional[Iterable[SOCKET_OPTION]] = None,
     ) -> None:
         self._connection: ConnectionInterface = HTTPConnection(
             origin=proxy_origin,
             keepalive_expiry=keepalive_expiry,
             network_backend=network_backend,
+            socket_options=socket_options,
         )
         self._proxy_origin = proxy_origin
         self._remote_origin = remote_origin

httpcore/backends/asyncio.py

@@ -13,7 +13,7 @@
     map_exceptions,
 )
 from .._utils import is_socket_readable
-from .base import AsyncNetworkBackend, AsyncNetworkStream
+from .base import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
 
 
 class AsyncIOStream(AsyncNetworkStream):
@@ -100,7 +100,10 @@ async def connect_tcp(
         port: int,
         timeout: typing.Optional[float] = None,
         local_address: typing.Optional[str] = None,
+        socket_options: typing.Optional[typing.Iterable[SOCKET_OPTION]] = None,
     ) -> AsyncNetworkStream:
+        if socket_options is None:
+            socket_options = []  # pragma: no cover
         exc_map = {
             TimeoutError: ConnectTimeout,
             OSError: ConnectError,
@@ -113,11 +116,18 @@ async def connect_tcp(
                     remote_port=port,
                     local_host=local_address,
                 )
+                for option in socket_options:
+                    stream._raw_socket.setsockopt(*option)  # type: ignore[attr-defined] # pragma: no cover
         return AsyncIOStream(stream)
 
     async def connect_unix_socket(
-        self, path: str, timeout: typing.Optional[float] = None
+        self,
+        path: str,
+        timeout: typing.Optional[float] = None,
+        socket_options: typing.Optional[typing.Iterable[SOCKET_OPTION]] = None,
     ) -> AsyncNetworkStream:  # pragma: nocover
+        if socket_options is None:
+            socket_options = []
         exc_map = {
             TimeoutError: ConnectTimeout,
             OSError: ConnectError,
@@ -126,6 +136,8 @@ async def connect_unix_socket(
         with map_exceptions(exc_map):
             with anyio.fail_after(timeout):
                 stream: anyio.abc.ByteStream = await anyio.connect_unix(path)
+                for option in socket_options:
+                    stream._raw_socket.setsockopt(*option)  # type: ignore[attr-defined] # pragma: no cover
         return AsyncIOStream(stream)
 
     async def sleep(self, seconds: float) -> None:

httpcore/backends/auto.py

@@ -1,8 +1,9 @@
+import typing
 from typing import Optional
 
 import sniffio
 
-from .base import AsyncNetworkBackend, AsyncNetworkStream
+from .base import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
 
 
 class AutoBackend(AsyncNetworkBackend):
@@ -24,17 +25,27 @@ async def connect_tcp(
         port: int,
         timeout: Optional[float] = None,
         local_address: Optional[str] = None,
+        socket_options: typing.Optional[typing.Iterable[SOCKET_OPTION]] = None,
     ) -> AsyncNetworkStream:
         await self._init_backend()
         return await self._backend.connect_tcp(
-            host, port, timeout=timeout, local_address=local_address
+            host,
+            port,
+            timeout=timeout,
+            local_address=local_address,
+            socket_options=socket_options,
         )
 
     async def connect_unix_socket(
-        self, path: str, timeout: Optional[float] = None
+        self,
+        path: str,
+        timeout: Optional[float] = None,
+        socket_options: typing.Optional[typing.Iterable[SOCKET_OPTION]] = None,
     ) -> AsyncNetworkStream:  # pragma: nocover
         await self._init_backend()
-        return await self._backend.connect_unix_socket(path, timeout=timeout)
+        return await self._backend.connect_unix_socket(
+            path, timeout=timeout, socket_options=socket_options
+        )
 
     async def sleep(self, seconds: float) -> None:  # pragma: nocover
         await self._init_backend()