Upload folder using huggingface_hub (part 7)
c12ef43 verified 3 days ago
63.3 kB
from __future__ import annotations

import base64
import concurrent.futures
import contextlib
import json
import queue
import secrets
import sys
import threading
import typing
from types import TracebackType

if sys.version_info >= (3, 13):
    from queue import ShutDown as QueueShutDown  # pragma: no cover
    from typing import TypeVar  # pragma: no cover
else:
    from typing_extensions import TypeVar  # pragma: no cover

    class QueueShutDown(Exception):  # pragma: no cover
        """Stand-in for `queue.ShutDown`, which is never raised before Python 3.13."""


import anyio
import wsproto
import wsproto.utilities
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from wsproto.frame_protocol import CloseReason

from .._client import USE_CLIENT_DEFAULT
from .._config import (
    DEFAULT_KEEPALIVE_PING_INTERVAL_SECONDS,
    DEFAULT_KEEPALIVE_PING_TIMEOUT_SECONDS,
    DEFAULT_MAX_MESSAGE_SIZE_BYTES,
    DEFAULT_QUEUE_SIZE,
)
from .._models import Headers
from ._exceptions import (
    HTTPXWSException,
    WebSocketDisconnect,
    WebSocketInvalidTypeReceived,
    WebSocketNetworkError,
    WebSocketUpgradeError,
)
from ._ping import AsyncPingManager, PingManager
from ._transport import ASGIWebSocketAsyncNetworkStream

if typing.TYPE_CHECKING:
    from httpcore2 import AsyncNetworkStream, NetworkStream

    from .._client import AsyncClient, Client, UseClientDefault
    from .._models import Response
    from .._types import (
        AuthTypes,
        CookieTypes,
        HeaderTypes,
        QueryParamTypes,
        RequestExtensions,
        TimeoutTypes,
    )

JSONMode = typing.Literal["text", "binary"]
TaskFunction = typing.TypeVar("TaskFunction")
TaskResult = typing.TypeVar("TaskResult")
SyncSession = TypeVar("SyncSession", bound="WebSocketSession", default="WebSocketSession")
AsyncSession = TypeVar("AsyncSession", bound="AsyncWebSocketSession", default="AsyncWebSocketSession")


class ShouldClose(Exception):
    pass


class EndOfStream(Exception):
    pass


class WebSocketSession:
    """
    Sync context manager representing an opened WebSocket session.
    Attributes:
        subprotocol (typing.Optional[str]):
            Optional protocol that has been accepted by the server.
        response (Response | None):
            The webSocket handshake response.
    """

    subprotocol: str | None
    response: Response | None

    def __init__(
        self,
        stream: NetworkStream,
        *,
        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,
        response: Response | None = None,
    ) -> None:
        self.stream = stream
        self.connection = wsproto.connection.Connection(wsproto.ConnectionType.CLIENT)
        self.response = response
        if self.response is not None:
            self.subprotocol = self.response.headers.get("sec-websocket-protocol")
        else:
            self.subprotocol = None

        self._events: queue.Queue[wsproto.events.Event | HTTPXWSException] = queue.Queue(queue_size)

        self._ping_manager = PingManager()
        self._should_close = threading.Event()
        self._write_lock = threading.Lock()
        self._should_close_task: concurrent.futures.Future[bool] | None = None
        self._executor: concurrent.futures.ThreadPoolExecutor | None = None

        self._max_message_size_bytes = max_message_size_bytes
        self._queue_size = queue_size
        self._keepalive_ping_interval_seconds = keepalive_ping_interval_seconds
        self._keepalive_ping_timeout_seconds = keepalive_ping_timeout_seconds

    def _get_executor_should_close_task(
        self,
    ) -> tuple[concurrent.futures.ThreadPoolExecutor, concurrent.futures.Future[bool]]:
        if self._should_close_task is None:
            self._executor = concurrent.futures.ThreadPoolExecutor()
            self._should_close_task = self._executor.submit(self._should_close.wait)
        assert self._executor is not None
        return self._executor, self._should_close_task

    def __enter__(self) -> WebSocketSession:
        self._background_receive_task = threading.Thread(
            target=self._background_receive, args=(self._max_message_size_bytes,)
        )
        self._background_receive_task.start()

        self._background_keepalive_ping_task: threading.Thread | None = None
        if self._keepalive_ping_interval_seconds is not None:
            self._background_keepalive_ping_task = threading.Thread(
                target=self._background_keepalive_ping,
                args=(
                    self._keepalive_ping_interval_seconds,
                    self._keepalive_ping_timeout_seconds,
                ),
            )
            self._background_keepalive_ping_task.start()

        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        self.close()
        # A task blocked on a full events queue only exits once there's room, or the queue shuts down.
        if sys.version_info >= (3, 13):  # pragma: no cover - Python 3.13+ only
            self._events.shutdown(immediate=True)
            self._background_receive_task.join()
            if self._background_keepalive_ping_task is not None:
                self._background_keepalive_ping_task.join()
        else:  # pragma: no cover - Python < 3.13 only
            for task in (self._background_receive_task, self._background_keepalive_ping_task):
                while task is not None and task.is_alive():
                    with contextlib.suppress(queue.Empty):
                        self._events.get(timeout=0.01)

    def ping(self, payload: bytes = b"") -> threading.Event:
        """
        Send a Ping message.
        Args:
            payload:
                Payload to attach to the Ping event.
                Internally, it's used to track this specific event.
                If left empty, a random one will be generated.
        Returns:
            An event that can be used to wait for the corresponding Pong response.
        Examples:
            Send a Ping and wait for the Pong
                pong_callback = ws.ping()
                # Will block until the corresponding Pong is received.
                pong_callback.wait()
        """
        ping_id, callback = self._ping_manager.create(payload)
        event = wsproto.events.Ping(ping_id)
        self.send(event)
        return callback

    def send(self, event: wsproto.events.Event) -> None:
        """
        Send an Event message.
        Mainly useful to send events that are not supported by the library.
        Most of the time, [ping()][httpx_ws.WebSocketSession.ping],
        [send_text()][httpx_ws.WebSocketSession.send_text],
        [send_bytes()][httpx_ws.WebSocketSession.send_bytes]
        and [send_json()][httpx_ws.WebSocketSession.send_json] are preferred.
        Args:
            event: The event to send.
        Raises:
            WebSocketNetworkError: A network error occured.
        Examples:
            Send an event.
                event = wsproto.events.Message(b"Hello!")
                ws.send(event)
        """
        import httpcore2

        try:
            data = self.connection.send(event)
            with self._write_lock:
                self.stream.write(data)
        except httpcore2.WriteError as e:
            self.close(CloseReason.INTERNAL_ERROR, "Stream write error")
            raise WebSocketNetworkError() from e

    def send_text(self, data: str) -> None:
        """
        Send a text message.
        Args:
            data: The text to send.
        Raises:
            WebSocketNetworkError: A network error occured.
        Examples:
            Send a text message.
                ws.send_text("Hello!")
        """
        event = wsproto.events.TextMessage(data=data)
        self.send(event)

    def send_bytes(self, data: bytes) -> None:
        """
        Send a bytes message.
        Args:
            data: The data to send.
        Raises:
            WebSocketNetworkError: A network error occured.
        Examples:
            Send a bytes message.
                ws.send_bytes(b"Hello!")
        """
        event = wsproto.events.BytesMessage(data=data)
        self.send(event)

    def send_json(self, data: typing.Any, mode: JSONMode = "text") -> None:
        """
        Send JSON data.
        Args:
            data:
                The data to send. Must be serializable by [json.dumps][json.dumps].
            mode:
                The sending mode. Should either be `'text'` or `'bytes'`.
        Raises:
            WebSocketNetworkError: A network error occured.
        Examples:
            Send JSON data.
                data = {"message": "Hello!"}
                ws.send_json(data)
        """
        assert mode in ["text", "binary"]
        serialized_data = json.dumps(data)
        if mode == "text":
            self.send_text(serialized_data)
        else:
            self.send_bytes(serialized_data.encode("utf-8"))

    def receive(self, timeout: float | None = None) -> wsproto.events.Event:
        """
        Receive an event from the server.
        Mainly useful to receive raw [wsproto.events.Event][wsproto.events.Event].
        Most of the time, [receive_text()][httpx_ws.WebSocketSession.receive_text],
        [receive_bytes()][httpx_ws.WebSocketSession.receive_bytes],
        and [receive_json()][httpx_ws.WebSocketSession.receive_json] are preferred.
        Args:
            timeout:
                Number of seconds to wait for an event.
                If `None`, will block until an event is available.
        Returns:
            A raw [wsproto.events.Event][wsproto.events.Event].
        Raises:
            TimeoutError: No event was received before the timeout delay.
            WebSocketDisconnect: The server closed the websocket.
            WebSocketNetworkError: A network error occured.
        Examples:
            Wait for an event until one is available.
                try:
                    event = ws.receive()
                except WebSocketDisconnect:
                    print("Connection closed")
            Wait for an event for 2 seconds.
                try:
                    event = ws.receive(timeout=2.)
                except TimeoutError:
                    print("No event received.")
                except WebSocketDisconnect:
                    print("Connection closed")
        """
        try:
            event = self._events.get(block=True, timeout=timeout)
        except queue.Empty as e:
            raise TimeoutError from e
        except QueueShutDown as e:  # pragma: no cover - Python 3.13+ only
            raise WebSocketDisconnect() from e
        if isinstance(event, HTTPXWSException):
            raise event
        if isinstance(event, wsproto.events.CloseConnection):
            raise WebSocketDisconnect(event.code, event.reason)
        return event

    def receive_text(self, timeout: float | None = None) -> str:
        """
        Receive text from the server.
        Args:
            timeout:
                Number of seconds to wait for an event.
                If `None`, will block until an event is available.
        Returns:
            Text data.
        Raises:
            TimeoutError: No event was received before the timeout delay.
            WebSocketDisconnect: The server closed the websocket.
            WebSocketNetworkError: A network error occured.
            WebSocketInvalidTypeReceived: The received event was not a text message.
        Examples:
            Wait for text until available.
                try:
                    text = ws.receive_text()
                except WebSocketDisconnect:
                    print("Connection closed")
            Wait for text for 2 seconds.
                try:
                    event = ws.receive_text(timeout=2.)
                except TimeoutError:
                    print("No text received.")
                except WebSocketDisconnect:
                    print("Connection closed")
        """
        event = self.receive(timeout)
        if isinstance(event, wsproto.events.TextMessage):
            return event.data
        raise WebSocketInvalidTypeReceived(event)

    def receive_bytes(self, timeout: float | None = None) -> bytes:
        """
        Receive bytes from the server.
        Args:
            timeout:
                Number of seconds to wait for an event.
                If `None`, will block until an event is available.
        Returns:
            Bytes data.
        Raises:
            TimeoutError: No event was received before the timeout delay.
            WebSocketDisconnect: The server closed the websocket.
            WebSocketNetworkError: A network error occured.
            WebSocketInvalidTypeReceived: The received event was not a bytes message.
        Examples:
            Wait for bytes until available.
                try:
                    data = ws.receive_bytes()
                except WebSocketDisconnect:
                    print("Connection closed")
            Wait for bytes for 2 seconds.
                try:
                    data = ws.receive_bytes(timeout=2.)
                except TimeoutError:
                    print("No data received.")
                except WebSocketDisconnect:
                    print("Connection closed")
        """
        event = self.receive(timeout)
        if isinstance(event, wsproto.events.BytesMessage):
            return bytes(event.data)
        raise WebSocketInvalidTypeReceived(event)

    def receive_json(self, timeout: float | None = None, mode: JSONMode = "text") -> typing.Any:
        """
        Receive JSON data from the server.
        The received data should be parseable by [json.loads][json.loads].
        Args:
            timeout:
                Number of seconds to wait for an event.
                If `None`, will block until an event is available.
            mode:
                Receive mode. Should either be `'text'` or `'bytes'`.
        Returns:
            Parsed JSON data.
        Raises:
            TimeoutError: No event was received before the timeout delay.
            WebSocketDisconnect: The server closed the websocket.
            WebSocketNetworkError: A network error occured.
            WebSocketInvalidTypeReceived: The received event
                didn't correspond to the specified mode.
        Examples:
            Wait for data until available.
                try:
                    data = ws.receive_json()
                except WebSocketDisconnect:
                    print("Connection closed")
            Wait for data for 2 seconds.
                try:
                    data = ws.receive_json(timeout=2.)
                except TimeoutError:
                    print("No data received.")
                except WebSocketDisconnect:
                    print("Connection closed")
        """
        assert mode in ["text", "binary"]
        data: str | bytes
        if mode == "text":
            data = self.receive_text(timeout)
        elif mode == "binary":
            data = self.receive_bytes(timeout)
        return json.loads(data)

    def close(self, code: int = 1000, reason: str | None = None) -> None:
        """
        Close the WebSocket session.
        Internally, it'll send the
        [CloseConnection][wsproto.events.CloseConnection] event.
        *This method is automatically called when exiting the context manager.*
        Args:
            code:
                The integer close code to indicate why the connection has closed.
            reason:
                Additional reasoning for why the connection has closed.
        Examples:
            Close the WebSocket session.
                ws.close()
        """
        import httpcore2

        self._should_close.set()
        if self._executor is not None:
            self._executor.shutdown(False)
        if self.connection.state not in {
            wsproto.connection.ConnectionState.LOCAL_CLOSING,
            wsproto.connection.ConnectionState.CLOSED,
        }:
            event = wsproto.events.CloseConnection(code, reason)
            data = self.connection.send(event)
            try:
                with self._write_lock:
                    self.stream.write(data)
            except httpcore2.WriteError:
                pass
        self.stream.close()

    def _background_receive(self, max_bytes: int) -> None:
        """
        Background thread listening for data from the server.
        Internally, it'll:
        * Answer to Ping events.
        * Acknowledge Pong events.
        * Put other events in the [_events][_events]
        queue that'll eventually be consumed by the user.
        Args:
            max_bytes: The maximum chunk size to read at each iteration.
        """
        import httpcore2

        partial_message_buffer: str | bytes | None = None
        partial_message_size = 0
        try:
            while not self._should_close.is_set():
                data = self._wait_until_closed(self._read_stream, max_bytes)
                self.connection.receive_data(data)
                for event in self.connection.events():
                    if isinstance(event, wsproto.events.Ping):
                        data = self.connection.send(event.response())
                        with self._write_lock:
                            self.stream.write(data)
                        continue
                    if isinstance(event, wsproto.events.Pong):
                        self._ping_manager.ack(event.payload)
                        continue
                    if isinstance(event, wsproto.events.CloseConnection):
                        self._should_close.set()
                    if isinstance(event, wsproto.events.Message):
                        partial_message_size += len(event.data.encode() if isinstance(event.data, str) else event.data)
                        if partial_message_size > max_bytes:
                            self.close(CloseReason.MESSAGE_TOO_BIG, "Message too big")
                            self._events.put(WebSocketDisconnect(CloseReason.MESSAGE_TOO_BIG, "Message too big"))
                            break
                        # Unfinished message: bufferize
                        if not event.message_finished:
                            if partial_message_buffer is None:
                                partial_message_buffer = event.data
                            else:
                                partial_message_buffer += event.data
                        # Finished message but no buffer: just emit the event
                        elif partial_message_buffer is None:
                            partial_message_size = 0
                            self._events.put(event)
                        # Finished message with buffer: emit the full event
                        else:
                            event_type = type(event)
                            full_message_event = event_type(partial_message_buffer + event.data)
                            partial_message_buffer = None
                            partial_message_size = 0
                            self._events.put(full_message_event)
                        continue
                    self._events.put(event)
        except (httpcore2.ReadError, httpcore2.WriteError, EndOfStream):
            self.close(CloseReason.INTERNAL_ERROR, "Stream error")
            with contextlib.suppress(QueueShutDown):
                self._events.put(WebSocketNetworkError())
        except (ShouldClose, QueueShutDown):
            pass

    def _background_keepalive_ping(self, interval_seconds: float, timeout_seconds: float | None = None) -> None:
        try:
            while not self._should_close.is_set():
                should_close = self._wait_until_closed(self._should_close.wait, interval_seconds)
                if should_close:  # pragma: no cover
                    raise ShouldClose()
                try:
                    pong_callback = self.ping()
                except wsproto.utilities.LocalProtocolError:
                    return
                if timeout_seconds is not None:
                    acknowledged = self._wait_until_closed(pong_callback.wait, timeout_seconds)
                    if not acknowledged:
                        self.close(CloseReason.INTERNAL_ERROR, "Keepalive ping timeout")
                        self._events.put(WebSocketNetworkError())
        except (ShouldClose, QueueShutDown):
            pass

    def _wait_until_closed(
        self, callable: typing.Callable[..., TaskResult], *args: typing.Any, **kwargs: typing.Any
    ) -> TaskResult:
        try:
            executor, should_close_task = self._get_executor_should_close_task()
            todo_task = executor.submit(callable, *args, **kwargs)
        except RuntimeError as e:
            raise ShouldClose() from e
        else:
            done, _ = concurrent.futures.wait(
                (todo_task, should_close_task),  # type: ignore[misc]
                return_when=concurrent.futures.FIRST_COMPLETED,
            )
            if should_close_task in done:
                raise ShouldClose()
            assert todo_task in done
            result = todo_task.result()
        return result

    def _read_stream(self, max_bytes: int) -> bytes:
        data = self.stream.read(max_bytes)
        if data == b"":
            raise EndOfStream()
        return data


class AsyncWebSocketSession(anyio.AsyncContextManagerMixin):
    """
    Async context manager representing an opened WebSocket session.
    Internally, this session uses an anyio task group to manage background tasks.
    As a result, exceptions that are not caught inside the context manager
    and propagate out of the `async with` block will be wrapped
    in an [ExceptionGroup][ExceptionGroup].
    To handle them, use the `except*` syntax:
        async with AsyncWebSocketSession(stream) as ws:
            try:
                data = await ws.receive_text()
            except WebSocketDisconnect:
                # Caught inside the context manager: plain exception.
                print("Connection closed")
        # If not caught inside:
        try:
            async with AsyncWebSocketSession(stream) as ws:
                data = await ws.receive_text()
        except* WebSocketDisconnect:
            # Propagated out of the context manager: wrapped in ExceptionGroup.
            print("Connection closed")
    Attributes:
        subprotocol (typing.Optional[str]):
            Optional protocol that has been accepted by the server.
        response (Response | None):
            The webSocket handshake response.
    """

    subprotocol: str | None
    response: Response | None
    _send_event: MemoryObjectSendStream[wsproto.events.Event | HTTPXWSException]
    _receive_event: MemoryObjectReceiveStream[wsproto.events.Event | HTTPXWSException]

    def __init__(
        self,
        stream: AsyncNetworkStream,
        *,
        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,
        response: Response | None = None,
    ) -> None:
        self.stream = stream
        self.connection = wsproto.connection.Connection(wsproto.ConnectionType.CLIENT)
        self.response = response
        if self.response is not None:
            self.subprotocol = self.response.headers.get("sec-websocket-protocol")
        else:
            self.subprotocol = None

        self._ping_manager = AsyncPingManager()
        self._should_close = anyio.Event()
        self._write_lock = anyio.Lock()

        self._max_message_size_bytes = max_message_size_bytes
        self._queue_size = queue_size

        # Always disable keepalive ping when emulating ASGI
        if isinstance(stream, ASGIWebSocketAsyncNetworkStream):
            self._keepalive_ping_interval_seconds = None
            self._keepalive_ping_timeout_seconds = None
        else:
            self._keepalive_ping_interval_seconds = keepalive_ping_interval_seconds
            self._keepalive_ping_timeout_seconds = keepalive_ping_timeout_seconds

    @contextlib.asynccontextmanager
    async def __asynccontextmanager__(self) -> typing.AsyncGenerator[AsyncWebSocketSession, None]:
        self._send_event, self._receive_event = anyio.create_memory_object_stream[
            wsproto.events.Event | HTTPXWSException
        ]()
        self._background_task_group = anyio.create_task_group()

        async with self._send_event, self._receive_event, self._background_task_group:
            self._background_task_group.start_soon(self._background_receive, self._max_message_size_bytes)
            if self._keepalive_ping_interval_seconds is not None:
                self._background_task_group.start_soon(
                    self._background_keepalive_ping,
                    self._keepalive_ping_interval_seconds,
                    self._keepalive_ping_timeout_seconds,
                )

            try:
                yield self
            finally:
                self._background_task_group.cancel_scope.cancel()
                with anyio.CancelScope(shield=True):
                    await self.close()

    async def ping(self, payload: bytes = b"") -> anyio.Event:
        """
        Send a Ping message.
        Args:
            payload:
                Payload to attach to the Ping event.
                Internally, it's used to track this specific event.
                If left empty, a random one will be generated.
        Returns:
            An event that can be used to wait for the corresponding Pong response.
        Examples:
            Send a Ping and wait for the Pong
                pong_callback = await ws.ping()
                # Will block until the corresponding Pong is received.
                await pong_callback.wait()
        """
        ping_id, callback = self._ping_manager.create(payload)
        event = wsproto.events.Ping(ping_id)
        await self.send(event)
        return callback

    async def send(self, event: wsproto.events.Event) -> None:
        """
        Send an Event message.
        Mainly useful to send events that are not supported by the library.
        Most of the time, [ping()][httpx_ws.AsyncWebSocketSession.ping],
        [send_text()][httpx_ws.AsyncWebSocketSession.send_text],
        [send_bytes()][httpx_ws.AsyncWebSocketSession.send_bytes]
        and [send_json()][httpx_ws.AsyncWebSocketSession.send_json] are preferred.
        Args:
            event: The event to send.
        Raises:
            WebSocketNetworkError: A network error occured.
        Note:
            Exceptions not caught inside the context manager will be
            wrapped in an [ExceptionGroup][ExceptionGroup]. Use `except*` to catch them
            outside the `async with` block.
        Examples:
            Send an event.
                event = await wsproto.events.Message(b"Hello!")
                ws.send(event)
        """
        import httpcore2

        try:
            data = self.connection.send(event)
            async with self._write_lock:
                await self.stream.write(data)
        except httpcore2.WriteError as e:
            await self.close(CloseReason.INTERNAL_ERROR, "Stream write error")
            raise WebSocketNetworkError() from e

    async def send_text(self, data: str) -> None:
        """
        Send a text message.
        Args:
            data: The text to send.
        Raises:
            WebSocketNetworkError: A network error occured.
        Note:
            Exceptions not caught inside the context manager will be
            wrapped in an [ExceptionGroup][ExceptionGroup]. Use `except*` to catch them
            outside the `async with` block.
        Examples:
            Send a text message.
                await ws.send_text("Hello!")
        """
        event = wsproto.events.TextMessage(data=data)
        await self.send(event)

    async def send_bytes(self, data: bytes) -> None:
        """
        Send a bytes message.
        Args:
            data: The data to send.
        Raises:
            WebSocketNetworkError: A network error occured.
        Note:
            Exceptions not caught inside the context manager will be
            wrapped in an [ExceptionGroup][ExceptionGroup]. Use `except*` to catch them
            outside the `async with` block.
        Examples:
            Send a bytes message.
                await ws.send_bytes(b"Hello!")
        """
        event = wsproto.events.BytesMessage(data=data)
        await self.send(event)

    async def send_json(self, data: typing.Any, mode: JSONMode = "text") -> None:
        """
        Send JSON data.
        Args:
            data:
                The data to send. Must be serializable by [json.dumps][json.dumps].
            mode:
                The sending mode. Should either be `'text'` or `'bytes'`.
        Raises:
            WebSocketNetworkError: A network error occured.
        Note:
            Exceptions not caught inside the context manager will be
            wrapped in an [ExceptionGroup][ExceptionGroup]. Use `except*` to catch them
            outside the `async with` block.
        Examples:
            Send JSON data.
                data = {"message": "Hello!"}
                await ws.send_json(data)
        """
        assert mode in ["text", "binary"]
        serialized_data = json.dumps(data)
        if mode == "text":
            await self.send_text(serialized_data)
        else:
            await self.send_bytes(serialized_data.encode("utf-8"))

    async def receive(self, timeout: float | None = None) -> wsproto.events.Event:
        """
        Receive an event from the server.
        Mainly useful to receive raw [wsproto.events.Event][wsproto.events.Event].
        Most of the time, [receive_text()][httpx_ws.AsyncWebSocketSession.receive_text],
        [receive_bytes()][httpx_ws.AsyncWebSocketSession.receive_bytes],
        and [receive_json()][httpx_ws.AsyncWebSocketSession.receive_json] are preferred.
        Args:
            timeout:
                Number of seconds to wait for an event.
                If `None`, will block until an event is available.
        Returns:
            A raw [wsproto.events.Event][wsproto.events.Event].
        Raises:
            TimeoutError: No event was received before the timeout delay.
            WebSocketDisconnect: The server closed the websocket.
            WebSocketNetworkError: A network error occured.
        Note:
            Exceptions not caught inside the context manager will be
            wrapped in an [ExceptionGroup][ExceptionGroup]. Use `except*` to catch them
            outside the `async with` block.
        Examples:
            Wait for an event until one is available.
                try:
                    event = await ws.receive()
                except WebSocketDisconnect:
                    print("Connection closed")
            Wait for an event for 2 seconds.
                try:
                    event = await ws.receive(timeout=2.)
                except TimeoutError:
                    print("No event received.")
                except WebSocketDisconnect:
                    print("Connection closed")
        """
        with anyio.fail_after(timeout):
            event = await self._receive_event.receive()
        if isinstance(event, HTTPXWSException):
            raise event
        if isinstance(event, wsproto.events.CloseConnection):
            raise WebSocketDisconnect(event.code, event.reason)
        return event

    async def receive_text(self, timeout: float | None = None) -> str:
        """
        Receive text from the server.
        Args:
            timeout:
                Number of seconds to wait for an event.
                If `None`, will block until an event is available.
        Returns:
            Text data.
        Raises:
            TimeoutError: No event was received before the timeout delay.
            WebSocketDisconnect: The server closed the websocket.
            WebSocketNetworkError: A network error occured.
            WebSocketInvalidTypeReceived: The received event was not a text message.
        Note:
            Exceptions not caught inside the context manager will be
            wrapped in an [ExceptionGroup][ExceptionGroup]. Use `except*` to catch them
            outside the `async with` block.
        Examples:
            Wait for text until available.
                try:
                    text = await ws.receive_text()
                except WebSocketDisconnect:
                    print("Connection closed")
            Wait for text for 2 seconds.
                try:
                    event = await ws.receive_text(timeout=2.)
                except TimeoutError:
                    print("No text received.")
                except WebSocketDisconnect:
                    print("Connection closed")
        """
        event = await self.receive(timeout)
        if isinstance(event, wsproto.events.TextMessage):
            return event.data
        raise WebSocketInvalidTypeReceived(event)

    async def receive_bytes(self, timeout: float | None = None) -> bytes:
        """
        Receive bytes from the server.
        Args:
            timeout:
                Number of seconds to wait for an event.
                If `None`, will block until an event is available.
        Returns:
            Bytes data.
        Raises:
            TimeoutError: No event was received before the timeout delay.
            WebSocketDisconnect: The server closed the websocket.
            WebSocketNetworkError: A network error occured.
            WebSocketInvalidTypeReceived: The received event was not a bytes message.
        Note:
            Exceptions not caught inside the context manager will be
            wrapped in an [ExceptionGroup][ExceptionGroup]. Use `except*` to catch them
            outside the `async with` block.
        Examples:
            Wait for bytes until available.
                try:
                    data = await ws.receive_bytes()
                except WebSocketDisconnect:
                    print("Connection closed")
            Wait for bytes for 2 seconds.
                try:
                    data = await ws.receive_bytes(timeout=2.)
                except TimeoutError:
                    print("No data received.")
                except WebSocketDisconnect:
                    print("Connection closed")
        """
        event = await self.receive(timeout)
        if isinstance(event, wsproto.events.BytesMessage):
            return bytes(event.data)
        raise WebSocketInvalidTypeReceived(event)

    async def receive_json(self, timeout: float | None = None, mode: JSONMode = "text") -> typing.Any:
        """
        Receive JSON data from the server.
        The received data should be parseable by [json.loads][json.loads].
        Args:
            timeout:
                Number of seconds to wait for an event.
                If `None`, will block until an event is available.
            mode:
                Receive mode. Should either be `'text'` or `'bytes'`.
        Returns:
            Parsed JSON data.
        Raises:
            TimeoutError: No event was received before the timeout delay.
            WebSocketDisconnect: The server closed the websocket.
            WebSocketNetworkError: A network error occured.
            WebSocketInvalidTypeReceived: The received event
                didn't correspond to the specified mode.
        Note:
            Exceptions not caught inside the context manager will be
            wrapped in an [ExceptionGroup][ExceptionGroup]. Use `except*` to catch them
            outside the `async with` block.
        Examples:
            Wait for data until available.
                try:
                    data = await ws.receive_json()
                except WebSocketDisconnect:
                    print("Connection closed")
            Wait for data for 2 seconds.
                try:
                    data = await ws.receive_json(timeout=2.)
                except TimeoutError:
                    print("No data received.")
                except WebSocketDisconnect:
                    print("Connection closed")
        """
        assert mode in ["text", "binary"]
        data: str | bytes
        if mode == "text":
            data = await self.receive_text(timeout)
        elif mode == "binary":
            data = await self.receive_bytes(timeout)
        return json.loads(data)

    async def close(self, code: int = 1000, reason: str | None = None) -> None:
        """
        Close the WebSocket session.
        Internally, it'll send the
        [CloseConnection][wsproto.events.CloseConnection] event.
        *This method is automatically called when exiting the context manager.*
        Args:
            code:
                The integer close code to indicate why the connection has closed.
            reason:
                Additional reasoning for why the connection has closed.
        Examples:
            Close the WebSocket session.
                await ws.close()
        """
        import httpcore2

        self._should_close.set()
        if self.connection.state not in {
            wsproto.connection.ConnectionState.LOCAL_CLOSING,
            wsproto.connection.ConnectionState.CLOSED,
        }:
            event = wsproto.events.CloseConnection(code, reason)
            data = self.connection.send(event)
            try:
                async with self._write_lock:
                    await self.stream.write(data)
            except httpcore2.WriteError:
                pass
        await self.stream.aclose()

    async def _background_receive(self, max_bytes: int) -> None:
        """
        Background task listening for data from the server.
        Internally, it'll:
        * Answer to Ping events.
        * Acknowledge Pong events.
        * Put other events in the [_events][_events]
        queue that'll eventually be consumed by the user.
        Args:
            max_bytes: The maximum chunk size to read at each iteration.
        """
        import httpcore2

        partial_message_buffer: str | bytes | None = None
        partial_message_size = 0
        try:
            while not self._should_close.is_set():
                data = await self._read_stream(max_bytes)
                self.connection.receive_data(data)
                for event in self.connection.events():
                    if isinstance(event, wsproto.events.Ping):
                        data = self.connection.send(event.response())
                        async with self._write_lock:
                            await self.stream.write(data)
                        continue
                    if isinstance(event, wsproto.events.Pong):
                        self._ping_manager.ack(event.payload)
                        continue
                    if isinstance(event, wsproto.events.CloseConnection):
                        self._should_close.set()
                    if isinstance(event, wsproto.events.Message):
                        partial_message_size += len(event.data.encode() if isinstance(event.data, str) else event.data)
                        if partial_message_size > max_bytes:
                            await self.close(CloseReason.MESSAGE_TOO_BIG, "Message too big")
                            await self._send_event.send(
                                WebSocketDisconnect(CloseReason.MESSAGE_TOO_BIG, "Message too big")
                            )
                            break
                        # Unfinished message: bufferize
                        if not event.message_finished:
                            if partial_message_buffer is None:
                                partial_message_buffer = event.data
                            else:
                                partial_message_buffer += event.data
                        # Finished message but no buffer: just emit the event
                        elif partial_message_buffer is None:
                            partial_message_size = 0
                            await self._send_event.send(event)
                        # Finished message with buffer: emit the full event
                        else:
                            event_type = type(event)
                            full_message_event = event_type(partial_message_buffer + event.data)
                            partial_message_buffer = None
                            partial_message_size = 0
                            await self._send_event.send(full_message_event)
                        continue
                    await self._send_event.send(event)
        except (httpcore2.ReadError, httpcore2.WriteError, EndOfStream):
            await self.close(CloseReason.INTERNAL_ERROR, "Stream error")
            await self._send_event.send(WebSocketNetworkError())

    async def _background_keepalive_ping(self, interval_seconds: float, timeout_seconds: float | None = None) -> None:
        while not self._should_close.is_set():
            await anyio.sleep(interval_seconds)

            try:
                pong_callback = await self.ping()
            # Connection is closing, exit the task
            except wsproto.utilities.LocalProtocolError:
                return

            if timeout_seconds is not None:
                try:
                    with anyio.fail_after(timeout_seconds):
                        await pong_callback.wait()
                except TimeoutError:
                    await self.close(CloseReason.INTERNAL_ERROR, "Keepalive ping timeout")
                    await self._send_event.send(WebSocketNetworkError())

    async def _read_stream(self, max_bytes: int) -> bytes:
        data = await self.stream.read(max_bytes)
        if data == b"":
            raise EndOfStream()
        return data


def _get_headers(
    subprotocols: list[str] | None,
) -> dict[str, typing.Any]:
    headers = {
        "connection": "upgrade",
        "upgrade": "websocket",
        "sec-websocket-key": base64.b64encode(secrets.token_bytes(16)).decode("utf-8"),
        "sec-websocket-version": "13",
    }
    if subprotocols is not None:
        headers["sec-websocket-protocol"] = ", ".join(subprotocols)
    return headers


class WebSocketClient(typing.Generic[SyncSession]):
    """
    A sync WebSocket client.
    This class provides an API for connecting to WebSocket.
    Attributes:
        client:
            HTTPX client to use.
        max_message_size_bytes:
            Maximum incoming message size in bytes.
            Larger messages, including fragmented messages whose
            cumulative size exceeds the limit, close the connection
            with `MESSAGE_TOO_BIG`.
            Defaults to 65 KiB.
        queue_size:
            Size of the queue where the received messages will be held
            until they are consumed.
            If the queue is full, the client will stop receive messages
            from the server until the queue has room available.
            Defaults to 512.
        keepalive_ping_interval_seconds:
            Interval at which the client will automatically send a Ping event
            to keep the connection alive. Set it to `None` to disable this mechanism.
            Defaults to 20 seconds.
        keepalive_ping_timeout_seconds:
            Maximum delay the client will wait for an answer to its Ping event.
            If the delay is exceeded,
            [WebSocketNetworkError][httpx_ws.WebSocketNetworkError]
            will be raised and the connection closed.
            Defaults to 20 seconds.
        session_class:
            The session class to use.
            Defaults to [WebSocketSession][httpx_ws.WebSocketSession].
    """

    def __init__(
        self,
        client: Client,
        *,
        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,
        session_class: type[SyncSession] = WebSocketSession,  # type: ignore[assignment]
    ) -> None:
        self.client = client
        self.max_message_size_bytes = max_message_size_bytes
        self.queue_size = queue_size
        self.keepalive_ping_interval_seconds = keepalive_ping_interval_seconds
        self.keepalive_ping_timeout_seconds = keepalive_ping_timeout_seconds
        self.session_class = session_class

    @contextlib.contextmanager
    def connect(
        self,
        url: str,
        *,
        subprotocols: list[str] | None = None,
        params: QueryParamTypes | None = None,
        headers: HeaderTypes | None = None,
        cookies: CookieTypes | None = None,
        auth: AuthTypes | UseClientDefault | None = USE_CLIENT_DEFAULT,
        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
        extensions: RequestExtensions | None = None,
    ) -> typing.Generator[SyncSession, None, None]:
        """
        Start a sync WebSocket session.
        It returns a context manager that'll automatically
        call [close()][httpx_ws.WebSocketSession.close] when exiting.
        Args:
            url: The WebSocket URL.
            subprotocols:
                Optional list of subprotocols to negotiate with the server.
            params:
                Query parameters to include in the handshake request.
            headers:
                Headers to include in the handshake request.
            cookies:
                Cookies to include in the handshake request.
            auth:
                Authentication to use for the handshake request.
            follow_redirects:
                Whether to follow redirects on the handshake request.
            timeout:
                Timeout configuration for the handshake request.
            extensions:
                Request extensions for the handshake request.
        Returns:
            A [context manager][contextlib.AbstractContextManager]
                for [WebSocketSession][httpx_ws.WebSocketSession].
        Examples:
            Initialize the client and connect to a WebSocket.
                with httpx2.Client() as client:
                    ws_client = WebSocketClient(client)
                    with ws_client.connect("http://localhost:8000/ws") as ws:
                        message = ws.receive_text()
                        print(message)
                        ws.send_text("Hello!")
        """
        with self.client.stream(
            "GET",
            url,
            params=params,
            headers=Headers(headers) | _get_headers(subprotocols),
            cookies=cookies,
            auth=auth,
            follow_redirects=follow_redirects,
            timeout=timeout,
            extensions=extensions,
        ) as response:
            if response.status_code != 101:
                raise WebSocketUpgradeError(response)

            session = self.session_class(
                response.extensions["network_stream"],
                max_message_size_bytes=self.max_message_size_bytes,
                queue_size=self.queue_size,
                keepalive_ping_interval_seconds=self.keepalive_ping_interval_seconds,
                keepalive_ping_timeout_seconds=self.keepalive_ping_timeout_seconds,
                response=response,
            )
            with session:
                yield session


@contextlib.contextmanager
def connect_ws(
    url: str,
    client: Client | 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,
    subprotocols: list[str] | None = None,
    params: QueryParamTypes | None = None,
    headers: HeaderTypes | None = None,
    cookies: CookieTypes | None = None,
    auth: AuthTypes | UseClientDefault | None = USE_CLIENT_DEFAULT,
    follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
    timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
    extensions: RequestExtensions | None = None,
) -> typing.Generator[WebSocketSession, None, None]:
    """
    Start a sync WebSocket session.
    It returns a context manager that'll automatically
    call [close()][httpx_ws.WebSocketSession.close] when exiting.
    Args:
        url: The WebSocket URL.
        client:
            HTTPX client to use.
            If not provided, a default one will be initialized.
        max_message_size_bytes:
            Maximum incoming message size in bytes.
            Larger messages, including fragmented messages whose
            cumulative size exceeds the limit, close the connection
            with `MESSAGE_TOO_BIG`.
            Defaults to 65 KiB.
        queue_size:
            Size of the queue where the received messages will be held
            until they are consumed.
            If the queue is full, the client will stop receive messages
            from the server until the queue has room available.
            Defaults to 512.
        keepalive_ping_interval_seconds:
            Interval at which the client will automatically send a Ping event
            to keep the connection alive. Set it to `None` to disable this mechanism.
            Defaults to 20 seconds.
        keepalive_ping_timeout_seconds:
            Maximum delay the client will wait for an answer to its Ping event.
            If the delay is exceeded,
            [WebSocketNetworkError][httpx_ws.WebSocketNetworkError]
            will be raised and the connection closed.
            Defaults to 20 seconds.
        subprotocols:
            Optional list of subprotocols to negotiate with the server.
        params:
            Query parameters to include in the handshake request.
        headers:
            Headers to include in the handshake request.
        cookies:
            Cookies to include in the handshake request.
        auth:
            Authentication to use for the handshake request.
        follow_redirects:
            Whether to follow redirects on the handshake request.
        timeout:
            Timeout configuration for the handshake request.
        extensions:
            Request extensions for the handshake request.
    Returns:
        A [context manager][contextlib.AbstractContextManager]
            for [WebSocketSession][httpx_ws.WebSocketSession].
    Examples:
        Without explicit HTTPX client.
            with connect_ws("http://localhost:8000/ws") as ws:
                message = ws.receive_text()
                print(message)
                ws.send_text("Hello!")
        With explicit HTTPX client.
            with httpx2.Client() as client:
                with connect_ws("http://localhost:8000/ws", client) as ws:
                    message = ws.receive_text()
                    print(message)
                    ws.send_text("Hello!")
    """
    if client is None:
        from .._client import Client

        owned_client: contextlib.AbstractContextManager[Client] = Client()
    else:
        owned_client = contextlib.nullcontext(client)

    with owned_client as client:
        ws_client = WebSocketClient(
            client=client,
            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,
        )
        with ws_client.connect(
            url,
            subprotocols=subprotocols,
            params=params,
            headers=headers,
            cookies=cookies,
            auth=auth,
            follow_redirects=follow_redirects,
            timeout=timeout,
            extensions=extensions,
        ) as websocket:
            yield websocket


class AsyncWebSocketClient(typing.Generic[AsyncSession]):
    """
    An async WebSocket client.
    This class provides an API for connecting to WebSocket.
    Attributes:
        client:
            HTTPX client to use.
        max_message_size_bytes:
            Maximum incoming message size in bytes.
            Larger messages, including fragmented messages whose
            cumulative size exceeds the limit, close the connection
            with `MESSAGE_TOO_BIG`.
            Defaults to 65 KiB.
        queue_size:
            Size of the queue where the received messages will be held
            until they are consumed.
            If the queue is full, the client will stop receive messages
            from the server until the queue has room available.
            Defaults to 512.
        keepalive_ping_interval_seconds:
            Interval at which the client will automatically send a Ping event
            to keep the connection alive. Set it to `None` to disable this mechanism.
            Defaults to 20 seconds.
        keepalive_ping_timeout_seconds:
            Maximum delay the client will wait for an answer to its Ping event.
            If the delay is exceeded,
            [WebSocketNetworkError][httpx_ws.WebSocketNetworkError]
            will be raised in an [ExceptionGroup][ExceptionGroup] and the connection closed.
            Defaults to 20 seconds.
        session_class:
            The session class to use.
            Defaults to [AsyncWebSocketSession][httpx_ws.AsyncWebSocketSession].
    """

    def __init__(
        self,
        client: AsyncClient,
        *,
        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,
        session_class: type[AsyncSession] = AsyncWebSocketSession,  # type: ignore[assignment]
    ) -> None:
        self.client = client
        self.max_message_size_bytes = max_message_size_bytes
        self.queue_size = queue_size
        self.keepalive_ping_interval_seconds = keepalive_ping_interval_seconds
        self.keepalive_ping_timeout_seconds = keepalive_ping_timeout_seconds
        self.session_class = session_class

    @contextlib.asynccontextmanager
    async def connect(
        self,
        url: str,
        *,
        subprotocols: list[str] | None = None,
        params: QueryParamTypes | None = None,
        headers: HeaderTypes | None = None,
        cookies: CookieTypes | None = None,
        auth: AuthTypes | UseClientDefault | None = USE_CLIENT_DEFAULT,
        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
        extensions: RequestExtensions | None = None,
    ) -> typing.AsyncGenerator[AsyncSession, None]:
        """
        Start an async WebSocket session.
        It returns an async context manager that'll automatically
        call [close()][httpx_ws.AsyncWebSocketSession.close] when exiting.
        Args:
            url: The WebSocket URL.
            subprotocols:
                Optional list of subprotocols to negotiate with the server.
            params:
                Query parameters to include in the handshake request.
            headers:
                Headers to include in the handshake request.
            cookies:
                Cookies to include in the handshake request.
            auth:
                Authentication to use for the handshake request.
            follow_redirects:
                Whether to follow redirects on the handshake request.
            timeout:
                Timeout configuration for the handshake request.
            extensions:
                Request extensions for the handshake request.
        Returns:
            An [async context manager][contextlib.AbstractAsyncContextManager]
                for [AsyncWebSocketSession][httpx_ws.AsyncWebSocketSession].
        Examples:
            Initialize the client and connect to a WebSocket.
                async with httpx2.AsyncClient() as client:
                    ws_client = AsyncWebSocketClient(client)
                    async with ws_client.connect("http://localhost:8000/ws") as ws:
                        message = await ws.receive_text()
                        print(message)
                        await ws.send_text("Hello!")
        """
        async with self.client.stream(
            "GET",
            url,
            params=params,
            headers=Headers(headers) | _get_headers(subprotocols),
            cookies=cookies,
            auth=auth,
            follow_redirects=follow_redirects,
            timeout=timeout,
            extensions=extensions,
        ) as response:
            if response.status_code != 101:
                raise WebSocketUpgradeError(response)

            session = self.session_class(
                response.extensions["network_stream"],
                max_message_size_bytes=self.max_message_size_bytes,
                queue_size=self.queue_size,
                keepalive_ping_interval_seconds=self.keepalive_ping_interval_seconds,
                keepalive_ping_timeout_seconds=self.keepalive_ping_timeout_seconds,
                response=response,
            )
            async with session:
                yield session


@contextlib.asynccontextmanager
async def aconnect_ws(
    url: str,
    client: AsyncClient | 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,
    subprotocols: list[str] | None = None,
    params: QueryParamTypes | None = None,
    headers: HeaderTypes | None = None,
    cookies: CookieTypes | None = None,
    auth: AuthTypes | UseClientDefault | None = USE_CLIENT_DEFAULT,
    follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
    timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
    extensions: RequestExtensions | None = None,
) -> typing.AsyncGenerator[AsyncWebSocketSession, None]:
    """
    Start an async WebSocket session.
    It returns an async context manager that'll automatically
    call [close()][httpx_ws.AsyncWebSocketSession.close] when exiting.
    Args:
        url: The WebSocket URL.
        client:
            HTTPX client to use.
            If not provided, a default one will be initialized.
        max_message_size_bytes:
            Maximum incoming message size in bytes.
            Larger messages, including fragmented messages whose
            cumulative size exceeds the limit, close the connection
            with `MESSAGE_TOO_BIG`.
            Defaults to 65 KiB.
        queue_size:
            Size of the queue where the received messages will be held
            until they are consumed.
            If the queue is full, the client will stop receive messages
            from the server until the queue has room available.
            Defaults to 512.
        keepalive_ping_interval_seconds:
            Interval at which the client will automatically send a Ping event
            to keep the connection alive. Set it to `None` to disable this mechanism.
            Defaults to 20 seconds.
        keepalive_ping_timeout_seconds:
            Maximum delay the client will wait for an answer to its Ping event.
            If the delay is exceeded,
            [WebSocketNetworkError][httpx_ws.WebSocketNetworkError]
            will be raised in an [ExceptionGroup][ExceptionGroup] and the connection closed.
            Defaults to 20 seconds.
        subprotocols:
            Optional list of subprotocols to negotiate with the server.
        params:
            Query parameters to include in the handshake request.
        headers:
            Headers to include in the handshake request.
        cookies:
            Cookies to include in the handshake request.
        auth:
            Authentication to use for the handshake request.
        follow_redirects:
            Whether to follow redirects on the handshake request.
        timeout:
            Timeout configuration for the handshake request.
        extensions:
            Request extensions for the handshake request.
    Returns:
        An [async context manager][contextlib.AbstractAsyncContextManager]
            for [AsyncWebSocketSession][httpx_ws.AsyncWebSocketSession].
    Examples:
        Without explicit HTTPX client.
            async with aconnect_ws("http://localhost:8000/ws") as ws:
                message = await ws.receive_text()
                print(message)
                await ws.send_text("Hello!")
        With explicit HTTPX client.
            async with httpx2.AsyncClient() as client:
                async with aconnect_ws("http://localhost:8000/ws", client) as ws:
                    message = await ws.receive_text()
                    print(message)
                    await ws.send_text("Hello!")
    """
    if client is None:
        from .._client import AsyncClient

        owned_client: contextlib.AbstractAsyncContextManager[AsyncClient] = AsyncClient()
    else:
        owned_client = contextlib.nullcontext(client)

    async with owned_client as client:
        ws_client = AsyncWebSocketClient(
            client=client,
            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,
        )
        async with ws_client.connect(
            url,
            subprotocols=subprotocols,
            params=params,
            headers=headers,
            cookies=cookies,
            auth=auth,
            follow_redirects=follow_redirects,
            timeout=timeout,
            extensions=extensions,
        ) as websocket:
            yield websocket