Strict bytes type

docs/reference/types.rst

@@ -5,6 +5,10 @@ Types
 
 .. autodata:: Data
 
+.. autodata:: BytesLike
+
+.. autodata:: DataLike
+
 .. autodata:: LoggerLike
 
 .. autodata:: StatusLike

src/websockets/asyncio/connection.py

@@ -19,10 +19,10 @@
     ConnectionClosedOK,
     ProtocolError,
 )
-from ..frames import DATA_OPCODES, BytesLike, CloseCode, Frame, Opcode
+from ..frames import DATA_OPCODES, CloseCode, Frame, Opcode
 from ..http11 import Request, Response
 from ..protocol import CLOSED, OPEN, Event, Protocol, State
-from ..typing import Data, LoggerLike, Subprotocol
+from ..typing import BytesLike, Data, DataLike, LoggerLike, Subprotocol
 from .compatibility import (
     TimeoutError,
     aiter,
@@ -402,7 +402,7 @@ async def recv_streaming(self, decode: bool | None = None) -> AsyncIterator[Data
 
     async def send(
         self,
-        message: Data | Iterable[Data] | AsyncIterable[Data],
+        message: DataLike | Iterable[DataLike] | AsyncIterable[DataLike],
         text: bool | None = None,
     ) -> None:
         """
@@ -657,7 +657,7 @@ async def wait_closed(self) -> None:
         """
         await asyncio.shield(self.connection_lost_waiter)
 
-    async def ping(self, data: Data | None = None) -> Awaitable[float]:
+    async def ping(self, data: DataLike | None = None) -> Awaitable[float]:
         """
         Send a Ping_.
 
@@ -710,7 +710,7 @@ async def ping(self, data: Data | None = None) -> Awaitable[float]:
             self.protocol.send_ping(data)
             return pong_waiter
 
-    async def pong(self, data: Data = b"") -> None:
+    async def pong(self, data: DataLike = b"") -> None:
         """
         Send a Pong_.
 
@@ -1134,7 +1134,7 @@ def eof_received(self) -> None:
 
 def broadcast(
     connections: Iterable[Connection],
-    message: Data,
+    message: DataLike,
     raise_exceptions: bool = False,
 ) -> None:
     """

src/websockets/asyncio/messages.py

@@ -179,6 +179,7 @@ async def get(self, decode: bool | None = None) -> Data:
         finally:
             self.get_in_progress = False
 
+        # This converts frame.data to bytes when it's a bytearray.
         data = b"".join(frame.data for frame in frames)
         if decode:
             return data.decode()
@@ -243,7 +244,8 @@ async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
             decoder = UTF8Decoder()
             yield decoder.decode(frame.data, frame.fin)
         else:
-            yield frame.data
+            # Convert to bytes when frame.data is a bytearray.
+            yield bytes(frame.data)
 
         # Following frames, for fragmented messages
         while not frame.fin:
@@ -257,7 +259,8 @@ async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
             if decode:
                 yield decoder.decode(frame.data, frame.fin)
             else:
-                yield frame.data
+                # Convert to bytes when frame.data is a bytearray.
+                yield bytes(frame.data)
 
         self.get_in_progress = False
 

src/websockets/extensions/permessage_deflate.py

@@ -13,7 +13,7 @@
     PayloadTooBig,
     ProtocolError,
 )
-from ..typing import ExtensionName, ExtensionParameter
+from ..typing import BytesLike, ExtensionName, ExtensionParameter
 from .base import ClientExtensionFactory, Extension, ServerExtensionFactory
 
 
@@ -129,6 +129,7 @@ def decode(
         # Uncompress data. Protect against zip bombs by preventing zlib from
         # decompressing more than max_length bytes (except when the limit is
         # disabled with max_size = None).
+        data: BytesLike
         if frame.fin and len(frame.data) < 2044:
             # Profiling shows that appending four bytes, which makes a copy, is
             # faster than calling decompress() again when data is less than 2kB.
@@ -182,6 +183,7 @@ def encode(self, frame: frames.Frame) -> frames.Frame:
                 )
 
         # Compress data.
+        data: BytesLike
         data = self.encoder.compress(frame.data) + self.encoder.flush(zlib.Z_SYNC_FLUSH)
         if frame.fin:
             # Sync flush generates between 5 or 6 bytes, ending with the bytes

src/websockets/frames.py

@@ -10,6 +10,7 @@
 from typing import Callable
 
 from .exceptions import PayloadTooBig, ProtocolError
+from .typing import BytesLike
 
 
 try:
@@ -118,9 +119,6 @@ class CloseCode(enum.IntEnum):
 }
 
 
-BytesLike = bytes, bytearray, memoryview
-
-
 @dataclasses.dataclass
 class Frame:
     """
@@ -140,7 +138,7 @@ class Frame:
     """
 
     opcode: Opcode
-    data: bytes | bytearray | memoryview
+    data: BytesLike
     fin: bool = True
     rsv1: bool = False
     rsv2: bool = False
@@ -202,7 +200,7 @@ def __str__(self) -> str:
     @classmethod
     def parse(
         cls,
-        read_exact: Callable[[int], Generator[None, None, bytes]],
+        read_exact: Callable[[int], Generator[None, None, bytes | bytearray]],
         *,
         mask: bool,
         max_size: int | None = None,
@@ -324,6 +322,7 @@ def serialize(
             output.write(mask_bytes)
 
         # Prepare the data.
+        data: BytesLike
         if mask:
             data = apply_mask(self.data, mask_bytes)
         else:
@@ -383,7 +382,7 @@ def __str__(self) -> str:
         return result
 
     @classmethod
-    def parse(cls, data: bytes) -> Close:
+    def parse(cls, data: BytesLike) -> Close:
         """
         Parse the payload of a close frame.
 
@@ -395,6 +394,8 @@ def parse(cls, data: bytes) -> Close:
             UnicodeDecodeError: If the reason isn't valid UTF-8.
 
         """
+        if isinstance(data, memoryview):
+            raise AssertionError("only compressed outgoing frames use memoryview")
         if len(data) >= 2:
             (code,) = struct.unpack("!H", data[:2])
             reason = data[2:].decode()

src/websockets/http11.py

@@ -47,7 +47,7 @@
 MAX_BODY_SIZE = int(os.environ.get("WEBSOCKETS_MAX_BODY_SIZE", "1_048_576"))  # 1 MiB
 
 
-def d(value: bytes) -> str:
+def d(value: bytes | bytearray) -> str:
     """
     Decode a bytestring for interpolating into an error message.
 
@@ -102,7 +102,7 @@ def exception(self) -> Exception | None:  # pragma: no cover
     @classmethod
     def parse(
         cls,
-        read_line: Callable[[int], Generator[None, None, bytes]],
+        read_line: Callable[[int], Generator[None, None, bytes | bytearray]],
     ) -> Generator[None, None, Request]:
         """
         Parse a WebSocket handshake request.
@@ -194,7 +194,7 @@ class Response:
     status_code: int
     reason_phrase: str
     headers: Headers
-    body: bytes = b""
+    body: bytes | bytearray = b""
 
     _exception: Exception | None = None
 
@@ -210,9 +210,9 @@ def exception(self) -> Exception | None:  # pragma: no cover
     @classmethod
     def parse(
         cls,
-        read_line: Callable[[int], Generator[None, None, bytes]],
-        read_exact: Callable[[int], Generator[None, None, bytes]],
-        read_to_eof: Callable[[int], Generator[None, None, bytes]],
+        read_line: Callable[[int], Generator[None, None, bytes | bytearray]],
+        read_exact: Callable[[int], Generator[None, None, bytes | bytearray]],
+        read_to_eof: Callable[[int], Generator[None, None, bytes | bytearray]],
         proxy: bool = False,
     ) -> Generator[None, None, Response]:
         """
@@ -276,6 +276,7 @@ def parse(
 
         headers = yield from parse_headers(read_line)
 
+        body: bytes | bytearray
         if proxy:
             body = b""
         else:
@@ -299,8 +300,8 @@ def serialize(self) -> bytes:
 
 
 def parse_line(
-    read_line: Callable[[int], Generator[None, None, bytes]],
-) -> Generator[None, None, bytes]:
+    read_line: Callable[[int], Generator[None, None, bytes | bytearray]],
+) -> Generator[None, None, bytes | bytearray]:
     """
     Parse a single line.
 
@@ -326,7 +327,7 @@ def parse_line(
 
 
 def parse_headers(
-    read_line: Callable[[int], Generator[None, None, bytes]],
+    read_line: Callable[[int], Generator[None, None, bytes | bytearray]],
 ) -> Generator[None, None, Headers]:
     """
     Parse HTTP headers.
@@ -379,10 +380,10 @@ def parse_headers(
 def read_body(
     status_code: int,
     headers: Headers,
-    read_line: Callable[[int], Generator[None, None, bytes]],
-    read_exact: Callable[[int], Generator[None, None, bytes]],
-    read_to_eof: Callable[[int], Generator[None, None, bytes]],
-) -> Generator[None, None, bytes]:
+    read_line: Callable[[int], Generator[None, None, bytes | bytearray]],
+    read_exact: Callable[[int], Generator[None, None, bytes | bytearray]],
+    read_to_eof: Callable[[int], Generator[None, None, bytes | bytearray]],
+) -> Generator[None, None, bytes | bytearray]:
     # https://datatracker.ietf.org/doc/html/rfc7230#section-3.3.3
 
     # Since websockets only does GET requests (no HEAD, no CONNECT), all

src/websockets/legacy/framing.py

@@ -6,8 +6,7 @@
 
 from .. import extensions, frames
 from ..exceptions import PayloadTooBig, ProtocolError
-from ..frames import BytesLike
-from ..typing import Data
+from ..typing import BytesLike, DataLike
 
 
 try:
@@ -19,7 +18,7 @@
 class Frame(NamedTuple):
     fin: bool
     opcode: frames.Opcode
-    data: bytes
+    data: BytesLike
     rsv1: bool = False
     rsv2: bool = False
     rsv3: bool = False
@@ -147,7 +146,7 @@ def write(
         write(self.new_frame.serialize(mask=mask, extensions=extensions))
 
 
-def prepare_data(data: Data) -> tuple[int, bytes]:
+def prepare_data(data: DataLike) -> tuple[int, BytesLike]:
     """
     Convert a string or byte-like object to an opcode and a bytes-like object.
 
@@ -171,7 +170,7 @@ def prepare_data(data: Data) -> tuple[int, bytes]:
         raise TypeError("data must be str or bytes-like")
 
 
-def prepare_ctrl(data: Data) -> bytes:
+def prepare_ctrl(data: DataLike) -> bytes:
     """
     Convert a string or byte-like object to bytes.