Add decode_text parameter to WebSocket for receiving TEXT as bytes (#11764)

Co-authored-by: Sam Bull <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
(cherry picked from commit 8b919d3)

Author: 3 people

CHANGES/11763.feature.rst

@@ -0,0 +1 @@
+Added ``decode_text`` parameter to :meth:`~aiohttp.ClientSession.ws_connect` and :class:`~aiohttp.web.WebSocketResponse` to receive WebSocket TEXT messages as raw bytes instead of decoded strings, enabling direct use with high-performance JSON parsers like ``orjson`` -- by :user:`bdraco`.

CHANGES/11764.feature.rst

@@ -0,0 +1 @@
+11763.feature.rst

aiohttp/_websocket/models.py

@@ -62,6 +62,25 @@ def json(self, *, loads: Callable[[Any], Any] = json.loads) -> Any:
         return loads(self.data)
 
 
+class WSMessageTextBytes(NamedTuple):
+    """WebSocket TEXT message with raw bytes (no UTF-8 decoding)."""
+
+    type: WSMsgType
+    data: bytes
+    extra: str | None
+
+    def json(self, *, loads: Callable[[bytes], Any] = json.loads) -> Any:
+        """Return parsed JSON data."""
+        return loads(self.data)
+
+
+# Type aliases for message types based on decode_text setting
+# When decode_text=True, TEXT messages have str data (WSMessage)
+# When decode_text=False, TEXT messages have bytes data (WSMessageTextBytes)
+WSMessageDecodeText = WSMessage
+WSMessageNoDecodeText = WSMessage | WSMessageTextBytes
+
+
 # Constructing the tuple directly to avoid the overhead of
 # the lambda and arg processing since NamedTuples are constructed
 # with a run time built lambda

aiohttp/_websocket/reader_c.pxd

@@ -26,6 +26,7 @@ cdef object TUPLE_NEW
 
 cdef object WSMsgType
 cdef object WSMessage
+cdef object WSMessageTextBytes
 
 cdef object WS_MSG_TYPE_TEXT
 cdef object WS_MSG_TYPE_BINARY
@@ -60,6 +61,7 @@ cdef class WebSocketReader:
 
     cdef WebSocketDataQueue queue
     cdef unsigned int _max_msg_size
+    cdef bint _decode_text
 
     cdef Exception _exc
     cdef bytearray _partial

aiohttp/_websocket/reader_py.py

@@ -15,6 +15,7 @@
     WebSocketError,
     WSCloseCode,
     WSMessage,
+    WSMessageTextBytes,
     WSMsgType,
 )
 
@@ -67,7 +68,7 @@ def __init__(
         self._eof = False
         self._waiter: asyncio.Future[None] | None = None
         self._exception: BaseException | None = None
-        self._buffer: deque[tuple[WSMessage, int]] = deque()
+        self._buffer: deque[tuple[WSMessage | WSMessageTextBytes, int]] = deque()
         self._get_buffer = self._buffer.popleft
         self._put_buffer = self._buffer.append
 
@@ -100,7 +101,9 @@ def feed_eof(self) -> None:
         self._release_waiter()
         self._exception = None  # Break cyclic references
 
-    def feed_data(self, data: "WSMessage", size: "cython_int") -> None:
+    def feed_data(
+        self, data: "WSMessage | WSMessageTextBytes", size: "cython_int"
+    ) -> None:
         self._size += size
         self._put_buffer((data, size))
         self._release_waiter()
@@ -132,10 +135,15 @@ def _read_from_buffer(self) -> WSMessage:
 
 class WebSocketReader:
     def __init__(
-        self, queue: WebSocketDataQueue, max_msg_size: int, compress: bool = True
+        self,
+        queue: WebSocketDataQueue,
+        max_msg_size: int,
+        compress: bool = True,
+        decode_text: bool = True,
     ) -> None:
         self.queue = queue
         self._max_msg_size = max_msg_size
+        self._decode_text = decode_text
 
         self._exc: Exception | None = None
         self._partial = bytearray()
@@ -262,21 +270,30 @@ def _handle_frame(
                 payload_merged = bytes(assembled_payload)
 
             if opcode == OP_CODE_TEXT:
-                try:
-                    text = payload_merged.decode("utf-8")
-                except UnicodeDecodeError as exc:
-                    raise WebSocketError(
-                        WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
-                    ) from exc
-
-                # XXX: The Text and Binary messages here can be a performance
-                # bottleneck, so we use tuple.__new__ to improve performance.
-                # This is not type safe, but many tests should fail in
-                # test_client_ws_functional.py if this is wrong.
-                self.queue.feed_data(
-                    TUPLE_NEW(WSMessage, (WS_MSG_TYPE_TEXT, text, "")),
-                    len(payload_merged),
-                )
+                if self._decode_text:
+                    try:
+                        text = payload_merged.decode("utf-8")
+                    except UnicodeDecodeError as exc:
+                        raise WebSocketError(
+                            WSCloseCode.INVALID_TEXT, "Invalid UTF-8 text message"
+                        ) from exc
+
+                    # XXX: The Text and Binary messages here can be a performance
+                    # bottleneck, so we use tuple.__new__ to improve performance.
+                    # This is not type safe, but many tests should fail in
+                    # test_client_ws_functional.py if this is wrong.
+                    self.queue.feed_data(
+                        TUPLE_NEW(WSMessage, (WS_MSG_TYPE_TEXT, text, "")),
+                        len(payload_merged),
+                    )
+                else:
+                    # Return raw bytes for TEXT messages when decode_text=False
+                    self.queue.feed_data(
+                        TUPLE_NEW(
+                            WSMessageTextBytes, (WS_MSG_TYPE_TEXT, payload_merged, "")
+                        ),
+                        len(payload_merged),
+                    )
             else:
                 self.queue.feed_data(
                     TUPLE_NEW(WSMessage, (WS_MSG_TYPE_BINARY, payload_merged, "")),

aiohttp/client.py

@@ -11,6 +11,7 @@
 from collections.abc import (
     Awaitable,
     Callable,
+    Collection,
     Coroutine,
     Generator,
     Iterable,
@@ -19,7 +20,16 @@
 )
 from contextlib import suppress
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Final, Generic, TypedDict, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Final,
+    Generic,
+    Literal,
+    TypedDict,
+    TypeVar,
+    overload,
+)
 
 import attr
 from multidict import CIMultiDict, MultiDict, MultiDictProxy, istr
@@ -186,6 +196,30 @@ class _RequestOptions(TypedDict, total=False):
     middlewares: Sequence[ClientMiddlewareType] | None
 
 
+class _WSConnectOptions(TypedDict, total=False):
+    method: str
+    protocols: Collection[str]
+    timeout: "ClientWSTimeout | _SENTINEL"
+    receive_timeout: float | None
+    autoclose: bool
+    autoping: bool
+    heartbeat: float | None
+    auth: BasicAuth | None
+    origin: str | None
+    params: Query
+    headers: LooseHeaders | None
+    proxy: StrOrURL | None
+    proxy_auth: BasicAuth | None
+    ssl: SSLContext | bool | Fingerprint
+    verify_ssl: bool | None
+    fingerprint: bytes | None
+    ssl_context: SSLContext | None
+    server_hostname: str | None
+    proxy_headers: LooseHeaders | None
+    compress: int
+    max_msg_size: int
+
+
 @attr.s(auto_attribs=True, frozen=True, slots=True)
 class ClientTimeout:
     total: float | None = None
@@ -214,7 +248,11 @@ class ClientTimeout:
 # https://www.rfc-editor.org/rfc/rfc9110#section-9.2.2
 IDEMPOTENT_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "TRACE", "PUT", "DELETE"})
 
-_RetType = TypeVar("_RetType", ClientResponse, ClientWebSocketResponse)
+_RetType_co = TypeVar(
+    "_RetType_co",
+    bound="ClientResponse | ClientWebSocketResponse[bool]",
+    covariant=True,
+)
 _CharsetResolver = Callable[[ClientResponse, bytes], str]
 
 
@@ -917,6 +955,35 @@ async def _connect_and_send_request(
                 )
             raise
 
+    if sys.version_info >= (3, 11) and TYPE_CHECKING:
+
+        @overload
+        def ws_connect(
+            self,
+            url: StrOrURL,
+            *,
+            decode_text: Literal[True] = ...,
+            **kwargs: Unpack[_WSConnectOptions],
+        ) -> "_BaseRequestContextManager[ClientWebSocketResponse[Literal[True]]]": ...
+
+        @overload
+        def ws_connect(
+            self,
+            url: StrOrURL,
+            *,
+            decode_text: Literal[False],
+            **kwargs: Unpack[_WSConnectOptions],
+        ) -> "_BaseRequestContextManager[ClientWebSocketResponse[Literal[False]]]": ...
+
+        @overload
+        def ws_connect(
+            self,
+            url: StrOrURL,
+            *,
+            decode_text: bool = ...,
+            **kwargs: Unpack[_WSConnectOptions],
+        ) -> "_BaseRequestContextManager[ClientWebSocketResponse[bool]]": ...
+
     def ws_connect(
         self,
         url: StrOrURL,
@@ -942,7 +1009,8 @@ def ws_connect(
         proxy_headers: LooseHeaders | None = None,
         compress: int = 0,
         max_msg_size: int = 4 * 1024 * 1024,
-    ) -> "_WSRequestContextManager":
+        decode_text: bool = True,
+    ) -> "_BaseRequestContextManager[ClientWebSocketResponse[bool]]":
         """Initiate websocket connection."""
         return _WSRequestContextManager(
             self._ws_connect(
@@ -968,9 +1036,39 @@ def ws_connect(
                 proxy_headers=proxy_headers,
                 compress=compress,
                 max_msg_size=max_msg_size,
+                decode_text=decode_text,
             )
         )
 
+    if sys.version_info >= (3, 11) and TYPE_CHECKING:
+
+        @overload
+        async def _ws_connect(
+            self,
+            url: StrOrURL,
+            *,
+            decode_text: Literal[True] = ...,
+            **kwargs: Unpack[_WSConnectOptions],
+        ) -> "ClientWebSocketResponse[Literal[True]]": ...
+
+        @overload
+        async def _ws_connect(
+            self,
+            url: StrOrURL,
+            *,
+            decode_text: Literal[False],
+            **kwargs: Unpack[_WSConnectOptions],
+        ) -> "ClientWebSocketResponse[Literal[False]]": ...
+
+        @overload
+        async def _ws_connect(
+            self,
+            url: StrOrURL,
+            *,
+            decode_text: bool = ...,
+            **kwargs: Unpack[_WSConnectOptions],
+        ) -> "ClientWebSocketResponse[bool]": ...
+
     async def _ws_connect(
         self,
         url: StrOrURL,
@@ -996,7 +1094,8 @@ async def _ws_connect(
         proxy_headers: LooseHeaders | None = None,
         compress: int = 0,
         max_msg_size: int = 4 * 1024 * 1024,
-    ) -> ClientWebSocketResponse:
+        decode_text: bool = True,
+    ) -> "ClientWebSocketResponse[bool]":
         if timeout is not sentinel:
             if isinstance(timeout, ClientWSTimeout):
                 ws_timeout = timeout
@@ -1162,7 +1261,9 @@ async def _ws_connect(
             transport = conn.transport
             assert transport is not None
             reader = WebSocketDataQueue(conn_proto, 2**16, loop=self._loop)
-            conn_proto.set_parser(WebSocketReader(reader, max_msg_size), reader)
+            conn_proto.set_parser(
+                WebSocketReader(reader, max_msg_size, decode_text=decode_text), reader
+            )
             writer = WebSocketWriter(
                 conn_proto,
                 transport,
@@ -1467,32 +1568,34 @@ async def __aexit__(
         await self.close()
 
 
-class _BaseRequestContextManager(Coroutine[Any, Any, _RetType], Generic[_RetType]):
+class _BaseRequestContextManager(
+    Coroutine[Any, Any, _RetType_co], Generic[_RetType_co]
+):
 
     __slots__ = ("_coro", "_resp")
 
-    def __init__(self, coro: Coroutine["asyncio.Future[Any]", None, _RetType]) -> None:
-        self._coro: Coroutine[asyncio.Future[Any], None, _RetType] = coro
+    def __init__(self, coro: Coroutine[asyncio.Future[Any], None, _RetType_co]) -> None:
+        self._coro: Coroutine[asyncio.Future[Any], None, _RetType_co] = coro
 
-    def send(self, arg: None) -> "asyncio.Future[Any]":
+    def send(self, arg: None) -> asyncio.Future[Any]:
         return self._coro.send(arg)
 
-    def throw(self, *args: Any, **kwargs: Any) -> "asyncio.Future[Any]":
+    def throw(self, *args: Any, **kwargs: Any) -> asyncio.Future[Any]:
         return self._coro.throw(*args, **kwargs)
 
     def close(self) -> None:
         return self._coro.close()
 
-    def __await__(self) -> Generator[Any, None, _RetType]:
+    def __await__(self) -> Generator[Any, None, _RetType_co]:
         ret = self._coro.__await__()
         return ret
 
-    def __iter__(self) -> Generator[Any, None, _RetType]:
+    def __iter__(self) -> Generator[Any, None, _RetType_co]:
         return self.__await__()
 
-    async def __aenter__(self) -> _RetType:
-        self._resp: _RetType = await self._coro
-        return await self._resp.__aenter__()
+    async def __aenter__(self) -> _RetType_co:
+        self._resp: _RetType_co = await self._coro
+        return await self._resp.__aenter__()  # type: ignore[return-value]
 
     async def __aexit__(
         self,
@@ -1504,7 +1607,7 @@ async def __aexit__(
 
 
 _RequestContextManager = _BaseRequestContextManager[ClientResponse]
-_WSRequestContextManager = _BaseRequestContextManager[ClientWebSocketResponse]
+_WSRequestContextManager = _BaseRequestContextManager[ClientWebSocketResponse[bool]]
 
 
 class _SessionRequestContextManager:
@@ -1513,7 +1616,7 @@ class _SessionRequestContextManager:
 
     def __init__(
         self,
-        coro: Coroutine["asyncio.Future[Any]", None, ClientResponse],
+        coro: Coroutine[asyncio.Future[Any], None, ClientResponse],
         session: ClientSession,
     ) -> None:
         self._coro = coro