Add native WebSocket support by vendoring httpx-ws

pyproject.toml

@@ -14,20 +14,23 @@ httpx2 = { workspace = true }
 
 [dependency-groups]
 dev = [
-  "httpx2[brotli,cli,http2,socks,zstd]",
+  "httpx2[brotli,cli,http2,socks,ws,zstd]",
   "httpcore2[asyncio,trio,http2,socks]",
   # Tests
   "chardet==6.0.0.post1",
   "coverage[toml]==7.10.6",
   "cryptography==46.0.7",
+  "flaky>=3.8",
   "pytest>=9.0.3",
   "pytest-codspeed>=4.1.1",
   "pytest-httpbin==2.0.0",
   "pytest-trio==0.8.0",
+  "starlette>=0.49",
   "trio==0.31.0",
   "trio-typing==0.10.0",
   "trustme==1.2.1",
   "uvicorn>=0.35",
+  "websockets>=15",
   "werkzeug>=3.1.6",
   # Linting
   "mypy==1.17.1",

src/httpx2/httpx2/__init__.py

@@ -1,3 +1,5 @@
+import typing as _typing
+
 from .__version__ import __description__, __title__, __version__
 from ._api import *
 from ._auth import *
@@ -11,15 +13,28 @@
 from ._types import *
 from ._urls import *
 
+if _typing.TYPE_CHECKING:
+    from ._websockets._api import AsyncWebSocketSession, WebSocketSession
+    from ._websockets._exceptions import (
+        HTTPXWSException,
+        WebSocketDisconnect,
+        WebSocketInvalidTypeReceived,
+        WebSocketNetworkError,
+        WebSocketUpgradeError,
+    )
+    from ._websockets._transport import ASGIWebSocketTransport
+
 __all__ = [
     "__description__",
     "__title__",
     "__version__",
     "ASGITransport",
+    "ASGIWebSocketTransport",
     "AsyncBaseTransport",
     "AsyncByteStream",
     "AsyncClient",
     "AsyncHTTPTransport",
+    "AsyncWebSocketSession",
     "Auth",
     "BaseTransport",
     "BasicAuth",
@@ -42,6 +57,7 @@
     "HTTPError",
     "HTTPStatusError",
     "HTTPTransport",
+    "HTTPXWSException",
     "InvalidURL",
     "Limits",
     "LocalProtocolError",
@@ -78,20 +94,37 @@
     "UnsupportedProtocol",
     "URL",
     "USE_CLIENT_DEFAULT",
+    "websocket",
+    "WebSocketDisconnect",
+    "WebSocketInvalidTypeReceived",
+    "WebSocketNetworkError",
+    "WebSocketSession",
+    "WebSocketUpgradeError",
     "WriteError",
     "WriteTimeout",
     "WSGITransport",
 ]
 
 
+_WEBSOCKET_NAMES = {
+    "ASGIWebSocketTransport",
+    "AsyncWebSocketSession",
+    "HTTPXWSException",
+    "WebSocketDisconnect",
+    "WebSocketInvalidTypeReceived",
+    "WebSocketNetworkError",
+    "WebSocketSession",
+    "WebSocketUpgradeError",
+}
+
 __locals = locals()
 for __name in __all__:
-    if not __name.startswith("__"):
+    if not __name.startswith("__") and __name not in _WEBSOCKET_NAMES:
         setattr(__locals[__name], "__module__", "httpx2")  # noqa
 
 
-def __getattr__(name: str) -> object:  # pragma: no cover
-    if name == "main":
+def __getattr__(name: str) -> object:
+    if name == "main":  # pragma: no cover
         import warnings
 
         warnings.warn(
@@ -104,4 +137,13 @@ def __getattr__(name: str) -> object:  # pragma: no cover
 
         return main
 
+    if name in _WEBSOCKET_NAMES:
+        from . import _websockets
+        from ._websockets._defaults import WS_EXTRA_INSTALL_MESSAGE
+
+        try:
+            return getattr(_websockets, name)
+        except ImportError:
+            raise ImportError(WS_EXTRA_INSTALL_MESSAGE)
+
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/httpx2/httpx2/_api.py

@@ -19,10 +19,18 @@
     TimeoutTypes,
 )
 from ._urls import URL
+from ._websockets._defaults import (
+    DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS,
+    DEFAULT_KEEPALIVE_PING_TIMEOUT_SECONDS,
+    DEFAULT_MAX_MESSAGE_SIZE_BYTES,
+    DEFAULT_QUEUE_SIZE,
+)
 
 if typing.TYPE_CHECKING:
     import ssl  # pragma: no cover
 
+    from ._websockets._api import WebSocketSession
+
 
 __all__ = [
     "delete",
@@ -34,6 +42,7 @@
     "put",
     "request",
     "stream",
+    "websocket",
 ]
 
 
@@ -424,3 +433,57 @@ def delete(
         timeout=timeout,
         trust_env=trust_env,
     )
+
+
+@contextmanager
+def websocket(
+    url: URL | str,
+    *,
+    params: QueryParamTypes | None = None,
+    headers: HeaderTypes | None = None,
+    cookies: CookieTypes | None = None,
+    auth: AuthTypes | None = None,
+    proxy: ProxyTypes | None = None,
+    follow_redirects: bool = False,
+    timeout: TimeoutTypes = DEFAULT_TIMEOUT_CONFIG,
+    verify: ssl.SSLContext | str | bool = True,
+    trust_env: bool = True,
+    subprotocols: list[str] | None = None,
+    max_message_size_bytes: int = DEFAULT_MAX_MESSAGE_SIZE_BYTES,
+    queue_size: int = DEFAULT_QUEUE_SIZE,
+    keepalive_ping_interval_seconds: float | None = DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS,
+    keepalive_ping_timeout_seconds: float | None = DEFAULT_KEEPALIVE_PING_TIMEOUT_SECONDS,
+) -> Generator[WebSocketSession]:
+    """
+    Open a WebSocket session.
+
+    The session is closed automatically when exiting the context manager.
+
+    ```python
+    with httpx2.websocket("ws://localhost:8000/ws") as ws:
+        ws.send_text("Hello!")
+        message = ws.receive_text()
+    ```
+
+    **Parameters**: See `httpx2.request` and `httpx2.Client.websocket`.
+    """
+    with Client(
+        cookies=cookies,
+        proxy=proxy,
+        verify=verify,
+        timeout=timeout,
+        trust_env=trust_env,
+    ) as client:
+        with client.websocket(
+            url,
+            params=params,
+            headers=headers,
+            auth=auth,
+            follow_redirects=follow_redirects,
+            subprotocols=subprotocols,
+            max_message_size_bytes=max_message_size_bytes,
+            queue_size=queue_size,
+            keepalive_ping_interval_seconds=keepalive_ping_interval_seconds,
+            keepalive_ping_timeout_seconds=keepalive_ping_timeout_seconds,
+        ) as session:
+            yield session

src/httpx2/httpx2/_client.py

@@ -48,10 +48,19 @@
 )
 from ._urls import URL, QueryParams
 from ._utils import URLPattern, get_environment_proxies
+from ._websockets._defaults import (
+    DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS,
+    DEFAULT_KEEPALIVE_PING_TIMEOUT_SECONDS,
+    DEFAULT_MAX_MESSAGE_SIZE_BYTES,
+    DEFAULT_QUEUE_SIZE,
+    WS_EXTRA_INSTALL_MESSAGE,
+)
 
 if typing.TYPE_CHECKING:
     import ssl  # pragma: no cover
 
+    from ._websockets._api import AsyncWebSocketSession, WebSocketSession
+
 __all__ = ["USE_CLIENT_DEFAULT", "AsyncClient", "Client"]
 
 # The type annotation for @classmethod and context managers here follows PEP 484
@@ -845,6 +854,47 @@ def stream(
         finally:
             response.close()
 
+    @contextmanager
+    def websocket(
+        self,
+        url: URL | str,
+        *,
+        max_message_size_bytes: int = DEFAULT_MAX_MESSAGE_SIZE_BYTES,
+        queue_size: int = DEFAULT_QUEUE_SIZE,
+        keepalive_ping_interval_seconds: float | None = DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS,
+        keepalive_ping_timeout_seconds: float | None = DEFAULT_KEEPALIVE_PING_TIMEOUT_SECONDS,
+        subprotocols: list[str] | None = None,
+        **kwargs: typing.Any,
+    ) -> Generator[WebSocketSession]:
+        """
+        Open a WebSocket session, using this client's configuration.
+
+        The session is closed automatically when exiting the context manager.
+
+        ```python
+        with httpx2.Client() as client:
+            with client.websocket("ws://localhost:8000/ws") as ws:
+                ws.send_text("Hello!")
+                message = ws.receive_text()
+        ```
+        """
+        try:
+            from ._websockets._api import connect_ws
+        except ImportError:
+            raise ImportError(WS_EXTRA_INSTALL_MESSAGE)
+
+        with connect_ws(
+            str(url),
+            self,
+            max_message_size_bytes=max_message_size_bytes,
+            queue_size=queue_size,
+            keepalive_ping_interval_seconds=keepalive_ping_interval_seconds,
+            keepalive_ping_timeout_seconds=keepalive_ping_timeout_seconds,
+            subprotocols=subprotocols,
+            **kwargs,
+        ) as session:
+            yield session
+
     def send(
         self,
         request: Request,
@@ -1548,6 +1598,47 @@ async def stream(
         finally:
             await response.aclose()
 
+    @asynccontextmanager
+    async def websocket(
+        self,
+        url: URL | str,
+        *,
+        max_message_size_bytes: int = DEFAULT_MAX_MESSAGE_SIZE_BYTES,
+        queue_size: int = DEFAULT_QUEUE_SIZE,
+        keepalive_ping_interval_seconds: float | None = DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS,
+        keepalive_ping_timeout_seconds: float | None = DEFAULT_KEEPALIVE_PING_TIMEOUT_SECONDS,
+        subprotocols: list[str] | None = None,
+        **kwargs: typing.Any,
+    ) -> AsyncGenerator[AsyncWebSocketSession]:
+        """
+        Open a WebSocket session, using this client's configuration.
+
+        The session is closed automatically when exiting the context manager.
+
+        ```python
+        async with httpx2.AsyncClient() as client:
+            async with client.websocket("ws://localhost:8000/ws") as ws:
+                await ws.send_text("Hello!")
+                message = await ws.receive_text()
+        ```
+        """
+        try:
+            from ._websockets._api import aconnect_ws
+        except ImportError:
+            raise ImportError(WS_EXTRA_INSTALL_MESSAGE)
+
+        async with aconnect_ws(
+            str(url),
+            self,
+            max_message_size_bytes=max_message_size_bytes,
+            queue_size=queue_size,
+            keepalive_ping_interval_seconds=keepalive_ping_interval_seconds,
+            keepalive_ping_timeout_seconds=keepalive_ping_timeout_seconds,
+            subprotocols=subprotocols,
+            **kwargs,
+        ) as session:
+            yield session
+
     async def send(
         self,
         request: Request,

src/httpx2/httpx2/_websockets/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 François Voron
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.