improve type handling for bytearray vs. bytes

fixes #185

Author: Kriechi

CHANGELOG.rst

@@ -10,6 +10,7 @@ Release History
 - Add docs for `wsproto.Connection`
 - Add support for Python 3.12, 3.13, and 3.14.
 - Drop support for Python 3.7, 3.8, and 3.9.
+- Improve Python typing, specifically bytes vs. bytearray.
 - Various linting, styling, and packaging improvements.
 
 1.2.0 (2022-08-23)

src/wsproto/connection.py

@@ -191,7 +191,7 @@ def events(self) -> Generator[Event, None, None]:
                 elif frame.opcode is Opcode.BINARY:
                     assert isinstance(frame.payload, (bytes, bytearray))
                     yield BytesMessage(
-                        data=frame.payload,
+                        data=bytearray(frame.payload),
                         frame_finished=frame.frame_finished,
                         message_finished=frame.message_finished,
                     )

src/wsproto/events.py

@@ -193,7 +193,7 @@ def response(self) -> CloseConnection:
         return CloseConnection(code=self.code, reason=self.reason)
 
 
-T = TypeVar("T", bytes, str)
+T = TypeVar("T", bytes, bytearray, str)
 
 
 @dataclass(frozen=True)
@@ -231,7 +231,7 @@ class Message(Event, Generic[T]):
 @dataclass(frozen=True)
 class TextMessage(Message[str]):  # pylint: disable=unsubscriptable-object
     """
-    This event is fired when a data frame with TEXT payload is received.
+    Fired when a data frame with TEXT payload is received.
 
     Fields:
 
@@ -240,31 +240,27 @@ class TextMessage(Message[str]):  # pylint: disable=unsubscriptable-object
        The message data as string, This only represents a single chunk
        of data and not a full WebSocket message.  You need to buffer
        and reassemble these chunks to get the full message.
-
     """
 
-    # https://github.com/python/mypy/issues/5744
     data: str
 
 
 @dataclass(frozen=True)
-class BytesMessage(Message[bytes]):  # pylint: disable=unsubscriptable-object
+class BytesMessage(Message[bytearray]):  # pylint: disable=unsubscriptable-object
     """
-    This event is fired when a data frame with BINARY payload is
-    received.
+    Fired when a data frame with BINARY payload is received.
 
     Fields:
 
     .. attribute:: data
 
-       The message data as byte string, can be decoded as UTF-8 for
+       The message data as bytearray, can be decoded as UTF-8 for
        TEXT messages.  This only represents a single chunk of data and
        not a full WebSocket message.  You need to buffer and
        reassemble these chunks to get the full message.
     """
 
-    # https://github.com/python/mypy/issues/5744
-    data: bytes
+    data: bytearray
 
 
 @dataclass(frozen=True)

src/wsproto/frame_protocol.py

@@ -23,12 +23,13 @@
 
 
 class XorMaskerSimple:
-    def __init__(self, masking_key: bytes) -> None:
+    def __init__(self, masking_key: bytearray | bytes) -> None:
         self._masking_key = masking_key
 
-    def process(self, data: bytes) -> bytes:
+    def process(self, data: bytearray) -> bytearray:
+        data = bytearray(data)
         if data:
-            data_array = bytearray(data)
+            data_array = data
             a, b, c, d = (_XOR_TABLE[n] for n in self._masking_key)
             data_array[::4] = data_array[::4].translate(a)
             data_array[1::4] = data_array[1::4].translate(b)
@@ -42,12 +43,12 @@ def process(self, data: bytes) -> bytes:
                 self._masking_key[key_rotation:] + self._masking_key[:key_rotation]
             )
 
-            return bytes(data_array)
+            return data_array
         return data
 
 
 class XorMaskerNull:
-    def process(self, data: bytes) -> bytes:
+    def process(self, data: bytearray) -> bytearray:
         return data
 
 
@@ -271,7 +272,7 @@ def consume_at_most(self, nbytes: int) -> bytearray:
         self.bytes_used += len(data)
         return data
 
-    def consume_exactly(self, nbytes: int) -> bytes | None:
+    def consume_exactly(self, nbytes: int) -> bytearray | None:
         if len(self.buffer) - self.bytes_used < nbytes:
             return None
 
@@ -370,11 +371,11 @@ def process_buffer(self) -> Frame | None:
         payload = self.masker.process(payload)
 
         for extension in self.extensions:
-            payload_ = extension.frame_inbound_payload_data(self, payload)
+            payload_ = extension.frame_inbound_payload_data(self, bytes(payload))
             if isinstance(payload_, CloseReason):
                 msg = "error in extension"
                 raise ParseFailed(msg, payload_)
-            payload = payload_
+            payload = bytearray(payload_)
 
         if finished:
             final = bytearray()
@@ -387,7 +388,7 @@ def process_buffer(self) -> Frame | None:
                     final += result
             payload += final
 
-        frame = Frame(self.effective_opcode, payload, finished, self.header.fin)
+        frame = Frame(self.effective_opcode, bytes(payload), finished, self.header.fin)
 
         if finished:
             self.header = None
@@ -585,7 +586,7 @@ def received_frames(self) -> Generator[Frame, None, None]:
             else:
                 yield event
 
-    def close(self, code: int | None = None, reason: str | None = None) -> bytes:
+    def close(self, code: int | None = None, reason: str | None = None) -> bytearray:
         payload = bytearray()
         if code is CloseReason.NO_STATUS_RCVD:
             code = None
@@ -603,23 +604,23 @@ def close(self, code: int | None = None, reason: str | None = None) -> bytes:
 
         return self._serialize_frame(Opcode.CLOSE, payload)
 
-    def ping(self, payload: bytes = b"") -> bytes:
+    def ping(self, payload: bytes = b"") -> bytearray:
         return self._serialize_frame(Opcode.PING, payload)
 
-    def pong(self, payload: bytes = b"") -> bytes:
+    def pong(self, payload: bytes = b"") -> bytearray:
         return self._serialize_frame(Opcode.PONG, payload)
 
     def send_data(
         self, payload: bytes | bytearray | str = b"", fin: bool = True,
-    ) -> bytes:
+    ) -> bytearray:
         if isinstance(payload, (bytes, bytearray, memoryview)):
             opcode = Opcode.BINARY
         elif isinstance(payload, str):
             opcode = Opcode.TEXT
             payload = payload.encode("utf-8")
         else:
             msg = "Must provide bytes or text"
-            raise ValueError(msg)
+            raise TypeError(msg)
 
         if self._outbound_opcode is None:
             self._outbound_opcode = opcode
@@ -642,11 +643,13 @@ def _make_fin_rsv_opcode(self, fin: bool, rsv: RsvBits, opcode: Opcode) -> int:
         return fin_bits | rsv_bits | opcode_bits
 
     def _serialize_frame(
-        self, opcode: Opcode, payload: bytes = b"", fin: bool = True,
+        self, opcode: Opcode, payload: bytes | bytearray = b"", fin: bool = True,
     ) -> bytearray:
+        payload = bytearray(payload)
+
         rsv = RsvBits(False, False, False)
         for extension in reversed(self.extensions):
-            rsv, payload = extension.frame_outbound(self, opcode, rsv, payload, fin)
+            rsv, payload = extension.frame_outbound(self, opcode, rsv, bytes(payload), fin)
 
         fin_rsv_opcode = self._make_fin_rsv_opcode(fin, rsv, opcode)
 
@@ -688,8 +691,8 @@ def _serialize_frame(
             # authors of malicious applications from selecting the bytes that
             # appear on the wire."
             #   -- https://tools.ietf.org/html/rfc6455#section-5.3
-            masking_key = os.urandom(4)
+            masking_key = bytearray(os.urandom(4))
             masker = XorMaskerSimple(masking_key)
-            return bytearray(header + masking_key + masker.process(payload))
+            return bytearray(header + masking_key + masker.process(bytearray(payload)))
 
         return bytearray(header + payload)

tests/test_frame_protocol.py

@@ -1197,7 +1197,7 @@ def test_data_we_have_no_idea_what_to_do_with(self) -> None:
         proto = fp.FrameProtocol(client=False, extensions=[])
         payload: Dict[str, str] = dict()
 
-        with pytest.raises(ValueError):
+        with pytest.raises(TypeError):
             # Intentionally passing illegal type.
             proto.send_data(payload)  # type: ignore