Add an option to disable decoding of text frames.

Also support decoding binary frames.

Fix #1376.

Author: aaugustin

src/websockets/asyncio/connection.py

@@ -172,7 +172,7 @@ async def __aiter__(self) -> AsyncIterator[Data]:
         except ConnectionClosedOK:
             return
 
-    async def recv(self) -> Data:
+    async def recv(self, decode: bool | None = None) -> Data:
         """
         Receive the next message.
 
@@ -192,21 +192,34 @@ async def recv(self) -> Data:
         When the message is fragmented, :meth:`recv` waits until all fragments
         are received, reassembles them, and returns the whole message.
 
+        Args:
+            decode: Set this flag to override the default behavior of returning
+                :class:`str` or :class:`bytes`. See below for details.
+
         Returns:
             A string (:class:`str`) for a Text_ frame or a bytestring
             (:class:`bytes`) for a Binary_ frame.
 
             .. _Text: https://www.rfc-editor.org/rfc/rfc6455.html#section-5.6
             .. _Binary: https://www.rfc-editor.org/rfc/rfc6455.html#section-5.6
 
+            You may override this behavior with the ``decode`` argument:
+
+            * Set ``decode=False`` to disable UTF-8 decoding of Text_ frames
+              and return a bytestring (:class:`bytes`). This may be useful to
+              optimize performance when decoding isn't needed.
+            * Set ``decode=True`` to force UTF-8 decoding of Binary_ frames
+              and return a string (:class:`str`). This is useful for servers
+              that send binary frames instead of text frames.
+
         Raises:
             ConnectionClosed: When the connection is closed.
             RuntimeError: If two coroutines call :meth:`recv` or
                 :meth:`recv_streaming` concurrently.
 
         """
         try:
-            return await self.recv_messages.get()
+            return await self.recv_messages.get(decode)
         except EOFError:
             raise self.protocol.close_exc from self.recv_exc
         except RuntimeError:
@@ -215,7 +228,7 @@ async def recv(self) -> Data:
                 "is already running recv or recv_streaming"
             ) from None
 
-    async def recv_streaming(self) -> AsyncIterator[Data]:
+    async def recv_streaming(self, decode: bool | None = None) -> AsyncIterator[Data]:
         """
         Receive the next message frame by frame.
 
@@ -232,21 +245,34 @@ async def recv_streaming(self) -> AsyncIterator[Data]:
         iterator in a partially consumed state, making the connection unusable.
         Instead, you should close the connection with :meth:`close`.
 
+        Args:
+            decode: Set this flag to override the default behavior of returning
+                :class:`str` or :class:`bytes`. See below for details.
+
         Returns:
             An iterator of strings (:class:`str`) for a Text_ frame or
             bytestrings (:class:`bytes`) for a Binary_ frame.
 
             .. _Text: https://www.rfc-editor.org/rfc/rfc6455.html#section-5.6
             .. _Binary: https://www.rfc-editor.org/rfc/rfc6455.html#section-5.6
 
+            You may override this behavior with the ``decode`` argument:
+
+            * Set ``decode=False`` to disable UTF-8 decoding of Text_ frames
+              and return bytestrings (:class:`bytes`). This may be useful to
+              optimize performance when decoding isn't needed.
+            * Set ``decode=True`` to force UTF-8 decoding of Binary_ frames
+              and return strings (:class:`str`). This is useful for servers
+              that send binary frames instead of text frames.
+
         Raises:
             ConnectionClosed: When the connection is closed.
             RuntimeError: If two coroutines call :meth:`recv` or
                 :meth:`recv_streaming` concurrently.
 
         """
         try:
-            async for frame in self.recv_messages.get_iter():
+            async for frame in self.recv_messages.get_iter(decode):
                 yield frame
         except EOFError:
             raise self.protocol.close_exc from self.recv_exc

src/websockets/asyncio/messages.py

@@ -121,6 +121,11 @@ async def get(self, decode: bool | None = None) -> Data:
         received, then it reassembles the message and returns it. To receive
         messages frame by frame, use :meth:`get_iter` instead.
 
+        Args:
+            decode: :obj:`False` disables UTF-8 decoding of text frames and
+                returns :class:`bytes`. :obj:`True` forces UTF-8 decoding of
+                binary frames and returns :class:`str`.
+
         Raises:
             EOFError: If the stream of frames has ended.
             RuntimeError: If two coroutines run :meth:`get` or :meth:`get_iter`
@@ -183,6 +188,11 @@ async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
         This method only makes sense for fragmented messages. If messages aren't
         fragmented, use :meth:`get` instead.
 
+        Args:
+            decode: :obj:`False` disables UTF-8 decoding of text frames and
+                returns :class:`bytes`. :obj:`True` forces UTF-8 decoding of
+                binary frames and returns :class:`str`.
+
         Raises:
             EOFError: If the stream of frames has ended.
             RuntimeError: If two coroutines run :meth:`get` or :meth:`get_iter`

tests/asyncio/test_connection.py

@@ -167,6 +167,16 @@ async def test_recv_binary(self):
         await self.remote_connection.send(b"\x01\x02\xfe\xff")
         self.assertEqual(await self.connection.recv(), b"\x01\x02\xfe\xff")
 
+    async def test_recv_encoded_text(self):
+        """recv receives an UTF-8 encoded text message."""
+        await self.remote_connection.send("😀")
+        self.assertEqual(await self.connection.recv(decode=False), "😀".encode())
+
+    async def test_recv_decoded_binary(self):
+        """recv receives an UTF-8 decoded binary message."""
+        await self.remote_connection.send("😀".encode())
+        self.assertEqual(await self.connection.recv(decode=True), "😀")
+
     async def test_recv_fragmented_text(self):
         """recv receives a fragmented text message."""
         await self.remote_connection.send(["😀", "😀"])
@@ -271,6 +281,22 @@ async def test_recv_streaming_binary(self):
             [b"\x01\x02\xfe\xff"],
         )
 
+    async def test_recv_streaming_encoded_text(self):
+        """recv_streaming receives an UTF-8 encoded text message."""
+        await self.remote_connection.send("😀")
+        self.assertEqual(
+            await alist(self.connection.recv_streaming(decode=False)),
+            ["😀".encode()],
+        )
+
+    async def test_recv_streaming_decoded_binary(self):
+        """recv_streaming receives a UTF-8 decoded binary message."""
+        await self.remote_connection.send("😀".encode())
+        self.assertEqual(
+            await alist(self.connection.recv_streaming(decode=True)),
+            ["😀"],
+        )
+
     async def test_recv_streaming_fragmented_text(self):
         """recv_streaming receives a fragmented text message."""
         await self.remote_connection.send(["😀", "😀"])