Enable mypy strict-optional

This ensures that values are properly checked for None.

Author: bluetech

Diff for: setup.cfg

@@ -46,7 +46,7 @@ disallow_untyped_defs = True
 ; implicit_reexport = False
 no_implicit_optional = True
 strict_equality = True
-strict_optional = False
+strict_optional = True
 warn_redundant_casts = True
 # warn_return_any = True
 warn_unused_configs = True

Diff for: test/helpers.py

@@ -17,11 +17,13 @@ def __init__(
         self.accept_response = accept_response
 
     def offer(self) -> Union[bool, str]:
+        assert self.offer_response is not None
         return self.offer_response
 
     def finalize(self, offer: str) -> None:
         self.accepted_offer = offer
 
     def accept(self, offer: str) -> Union[bool, str]:
+        assert self.accept_response is not None
         self.offered = offer
         return self.accept_response

Diff for: test/test_extensions.py

@@ -13,28 +13,33 @@ def test_offer(self) -> None:
 
     def test_accept(self) -> None:
         ext = wpext.Extension()
-        assert ext.accept(None) is None
+        offer = "myext"
+        assert ext.accept(offer) is None
 
     def test_finalize(self) -> None:
         ext = wpext.Extension()
-        ext.finalize(None)
+        offer = "myext"
+        ext.finalize(offer)
 
     def test_frame_inbound_header(self) -> None:
         ext = wpext.Extension()
-        result = ext.frame_inbound_header(None, None, None, None)
+        result = ext.frame_inbound_header(None, None, None, None)  # type: ignore
         assert result == fp.RsvBits(False, False, False)
 
     def test_frame_inbound_payload_data(self) -> None:
         ext = wpext.Extension()
         data = b""
-        assert ext.frame_inbound_payload_data(None, data) == data
+        assert ext.frame_inbound_payload_data(None, data) == data  # type: ignore
 
     def test_frame_inbound_complete(self) -> None:
         ext = wpext.Extension()
-        assert ext.frame_inbound_complete(None, None) is None
+        assert ext.frame_inbound_complete(None, None) is None  # type: ignore
 
     def test_frame_outbound(self) -> None:
         ext = wpext.Extension()
         rsv = fp.RsvBits(True, True, True)
         data = b""
-        assert ext.frame_outbound(None, None, rsv, data, None) == (rsv, data)
+        assert ext.frame_outbound(None, None, rsv, data, None) == (  # type: ignore
+            rsv,
+            data,
+        )

Diff for: test/test_frame_protocol.py

@@ -917,7 +917,10 @@ def test_long_text_message(self) -> None:
         assert frame.payload == payload
 
     def _close_test(
-        self, code: Optional[int], reason: str = None, reason_bytes: bytes = None
+        self,
+        code: Optional[int],
+        reason: Optional[str] = None,
+        reason_bytes: Optional[bytes] = None,
     ) -> None:
         payload = b""
         if code:

Diff for: test/test_permessage_deflate.py

@@ -296,7 +296,7 @@ def test_client_inbound_compressed_multiple_data_frames(self, client: bool) -> N
         data += result4
 
         result5 = ext.frame_inbound_complete(proto, True)
-        assert not isinstance(result5, fp.CloseReason)
+        assert isinstance(result5, bytes)
         data += result5
 
         assert data == payload

Diff for: wsproto/__init__.py

@@ -64,7 +64,7 @@ def send(self, event: Event) -> bytes:
             data += self.connection.send(event)
         return data
 
-    def receive_data(self, data: bytes) -> None:
+    def receive_data(self, data: Optional[bytes]) -> None:
         """
         Feed network data into the connection instance.
 

Diff for: wsproto/connection.py

@@ -109,7 +109,7 @@ def send(self, event: Event) -> bytes:
             raise LocalProtocolError("Event {} cannot be sent.".format(event))
         return data
 
-    def receive_data(self, data: bytes) -> None:
+    def receive_data(self, data: Optional[bytes]) -> None:
         """
         Pass some received data to the connection for handling.
 

Diff for: wsproto/extensions.py

@@ -13,15 +13,15 @@
 
 
 class Extension:
-    name: Optional[str] = None
+    name: str
 
     def enabled(self) -> bool:
         return False
 
     def offer(self) -> Union[bool, str]:
         pass
 
-    def accept(self, offer: str) -> Union[bool, str]:
+    def accept(self, offer: str) -> Optional[Union[bool, str]]:
         pass
 
     def finalize(self, offer: str) -> None:
@@ -123,7 +123,7 @@ def finalize(self, offer: str) -> None:
 
         self._enabled = True
 
-    def _parse_params(self, params: str) -> Tuple[int, int]:
+    def _parse_params(self, params: str) -> Tuple[Optional[int], Optional[int]]:
         client_max_window_bits = None
         server_max_window_bits = None
 
@@ -198,6 +198,7 @@ def frame_inbound_payload_data(
     ) -> Union[bytes, CloseReason]:
         if not self._inbound_compressed or not self._inbound_is_compressible:
             return data
+        assert self._decompressor is not None
 
         try:
             return self._decompressor.decompress(bytes(data))
@@ -215,6 +216,7 @@ def frame_inbound_complete(
         if not self._inbound_compressed:
             self._inbound_compressed = None
             return None
+        assert self._decompressor is not None
 
         try:
             data = self._decompressor.decompress(b"\x00\x00\xff\xff")

Diff for: wsproto/frame_protocol.py

@@ -211,7 +211,7 @@ class Header(NamedTuple):
     rsv: RsvBits
     opcode: Opcode
     payload_len: int
-    masking_key: bytes
+    masking_key: Optional[bytes]
 
 
 class Frame(NamedTuple):
@@ -238,7 +238,7 @@ def _truncate_utf8(data: bytes, nbytes: int) -> bytes:
 
 
 class Buffer:
-    def __init__(self, initial_bytes: bytes = None) -> None:
+    def __init__(self, initial_bytes: Optional[bytes] = None) -> None:
         self.buffer = bytearray()
         self.bytes_used = 0
         if initial_bytes:
@@ -255,7 +255,7 @@ def consume_at_most(self, nbytes: int) -> bytes:
         self.bytes_used += len(data)
         return data
 
-    def consume_exactly(self, nbytes: int) -> bytes:
+    def consume_exactly(self, nbytes: int) -> Optional[bytes]:
         if len(self.buffer) - self.bytes_used < nbytes:
             return None
 
@@ -297,7 +297,10 @@ def process_frame(self, frame: Frame) -> Frame:
             data = frame.payload
         else:
             assert isinstance(frame.payload, (bytes, bytearray))
-            data = self.decode_payload(frame.payload, finished)
+            try:
+                data = self.decoder.decode(frame.payload, finished)
+            except UnicodeDecodeError as exc:
+                raise ParseFailed(str(exc), CloseReason.INVALID_FRAME_PAYLOAD_DATA)
 
         frame = Frame(self.opcode, data, frame.frame_finished, finished)
 
@@ -307,15 +310,11 @@ def process_frame(self, frame: Frame) -> Frame:
 
         return frame
 
-    def decode_payload(self, data: bytes, finished: bool) -> str:
-        try:
-            return self.decoder.decode(data, finished)
-        except UnicodeDecodeError as exc:
-            raise ParseFailed(str(exc), CloseReason.INVALID_FRAME_PAYLOAD_DATA)
-
 
 class FrameDecoder:
-    def __init__(self, client: bool, extensions: List["Extension"] = None) -> None:
+    def __init__(
+        self, client: bool, extensions: Optional[List["Extension"]] = None
+    ) -> None:
         self.client = client
         self.extensions = extensions or []
 
@@ -330,10 +329,14 @@ def __init__(self, client: bool, extensions: List["Extension"] = None) -> None:
     def receive_bytes(self, data: bytes) -> None:
         self.buffer.feed(data)
 
-    def process_buffer(self) -> Frame:
+    def process_buffer(self) -> Optional[Frame]:
         if not self.header:
             if not self.parse_header():
                 return None
+        # parse_header() sets these.
+        assert self.header is not None
+        assert self.masker is not None
+        assert self.effective_opcode is not None
 
         if len(self.buffer) < self.payload_required:
             return None
@@ -398,8 +401,8 @@ def parse_header(self) -> bool:
             raise ParseFailed("Invalid attempt to fragment control frame")
 
         has_mask = bool(data[1] & MASK_MASK)
-        payload_len = data[1] & PAYLOAD_LEN_MASK
-        payload_len = self.parse_extended_payload_length(opcode, payload_len)
+        payload_len_short = data[1] & PAYLOAD_LEN_MASK
+        payload_len = self.parse_extended_payload_length(opcode, payload_len_short)
         if payload_len is None:
             self.buffer.rollback()
             return False
@@ -429,7 +432,9 @@ def parse_header(self) -> bool:
         self.payload_consumed = 0
         return True
 
-    def parse_extended_payload_length(self, opcode: Opcode, payload_len: int) -> int:
+    def parse_extended_payload_length(
+        self, opcode: Opcode, payload_len: int
+    ) -> Optional[int]:
         if opcode.iscontrol() and payload_len > MAX_PAYLOAD_NORMAL:
             raise ParseFailed("Control frame with payload len > 125")
         if payload_len == PAYLOAD_LENGTH_TWO_BYTE:
@@ -518,7 +523,7 @@ def _process_close(self, frame: Frame) -> Frame:
 
         return Frame(frame.opcode, data, frame.frame_finished, frame.message_finished)
 
-    def _parse_more_gen(self) -> Generator[Frame, None, None]:
+    def _parse_more_gen(self) -> Generator[Optional[Frame], None, None]:
         # Consume as much as we can from self._buffer, yielding events, and
         # then yield None when we need more data. Or raise ParseFailed.
 

Diff for: wsproto/handshake.py

@@ -102,7 +102,7 @@ def send(self, event: Event) -> bytes:
             )
         return data
 
-    def receive_data(self, data: bytes) -> None:
+    def receive_data(self, data: Optional[bytes]) -> None:
         """Receive data from the remote.
 
         A list of events that the remote peer triggered by sending
@@ -243,6 +243,8 @@ def _process_connection_request(self, event: h11.Request) -> Request:  # noqa: M
         return self._initiating_request
 
     def _accept(self, event: AcceptConnection) -> bytes:
+        # _accept is always called after _process_connection_request.
+        assert self._initiating_request is not None
         request_headers = normed_header_dict(self._initiating_request.extra_headers)
 
         nonce = request_headers[b"sec-websocket-key"]
@@ -354,6 +356,10 @@ def _initiate_connection(self, request: Request) -> bytes:
     def _establish_client_connection(
         self, event: h11.InformationalResponse
     ) -> AcceptConnection:  # noqa: MC0001
+        # _establish_client_connection is always called after _initiate_connection.
+        assert self._initiating_request is not None
+        assert self._nonce is not None
+
         accept = None
         connection_tokens = None
         accepts: List[str] = []