Small cython cleanps to the frame helpers (#733)

aioesphomeapi/_frame_helper/base.pxd

@@ -39,6 +39,6 @@ cdef class APIFrameHelper:
     @cython.locals(end_of_frame_pos="unsigned int")
     cdef void _remove_from_buffer(self)
 
-    cpdef write_packets(self, list packets, bint debug_enabled)
+    cpdef void write_packets(self, list packets, bint debug_enabled)
 
-    cdef void _write_bytes(self, bytes data, bint debug_enabled)
+    cdef void _write_bytes(self, object data, bint debug_enabled)

aioesphomeapi/_frame_helper/base.py

@@ -22,6 +22,7 @@ SOCKET_ERRORS = (
 WRITE_EXCEPTIONS = (RuntimeError, ConnectionResetError, OSError)
 
 _int = int
+_bytes = bytes
 
 
 class APIFrameHelper:
@@ -196,7 +197,7 @@ class APIFrameHelper:
     def resume_writing(self) -> None:
         """Stub."""
 
-    def _write_bytes(self, data: bytes, debug_enabled: bool) -> None:
+    def _write_bytes(self, data: _bytes, debug_enabled: bool) -> None:
         """Write bytes to the socket."""
         if debug_enabled:
             _LOGGER.debug("%s: Sending frame: [%s]", self._log_name, data.hex())

aioesphomeapi/_frame_helper/noise.pxd

@@ -29,7 +29,7 @@ cdef class APINoiseFrameHelper(APIFrameHelper):
         msg_size_high="unsigned char",
         msg_size_low="unsigned char",
     )
-    cpdef data_received(self, object data)
+    cpdef void data_received(self, object data)
 
     @cython.locals(
         msg=bytes,
@@ -64,6 +64,6 @@ cdef class APINoiseFrameHelper(APIFrameHelper):
         frame=bytes,
         frame_len=cython.uint,
     )
-    cpdef write_packets(self, list packets, bint debug_enabled)
+    cpdef void write_packets(self, list packets, bint debug_enabled)
 
     cdef _error_on_incorrect_preamble(self, bytes msg)

aioesphomeapi/_frame_helper/noise.py

@@ -140,8 +140,6 @@ class APINoiseFrameHelper(APIFrameHelper):
             if (header := self._read(3)) is None:
                 return
             preamble = header[0]
-            msg_size_high = header[1]
-            msg_size_low = header[2]
             if preamble != 0x01:
                 self._handle_error_and_close(
                     ProtocolAPIError(
@@ -149,6 +147,8 @@ class APINoiseFrameHelper(APIFrameHelper):
                     )
                 )
                 return
+            msg_size_high = header[1]
+            msg_size_low = header[2]
             if (frame := self._read((msg_size_high << 8) | msg_size_low)) is None:
                 # The complete frame is not yet available, wait for more data
                 # to arrive before continuing, since callback_packet has not

aioesphomeapi/_frame_helper/plain_text.pxd

@@ -5,12 +5,14 @@ from .base cimport APIFrameHelper
 
 
 cdef object varuint_to_bytes
+cdef bytes EMPTY_PACKET
+cdef bint TYPE_CHECKING
 
 cpdef _varuint_to_bytes(cython.int value)
 
 cdef class APIPlaintextFrameHelper(APIFrameHelper):
 
-    cpdef data_received(self, object data)
+    cpdef void data_received(self, object data)
 
     cdef void _error_on_incorrect_preamble(self, int preamble)
 
@@ -20,4 +22,4 @@ cdef class APIPlaintextFrameHelper(APIFrameHelper):
         packet=tuple,
         type_=object
     )
-    cpdef write_packets(self, list packets, bint debug_enabled)
+    cpdef void write_packets(self, list packets, bint debug_enabled)

aioesphomeapi/_frame_helper/plain_text.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 
 import asyncio
 from functools import lru_cache
+from typing import TYPE_CHECKING
 
 from ..core import ProtocolAPIError, RequiresEncryptionAPIError
 from .base import APIFrameHelper
 
 _int = int
 
+EMPTY_PACKET = b""
+
 
 def _varuint_to_bytes(value: _int) -> bytes:
     """Convert a varuint to bytes."""
@@ -71,17 +74,19 @@ class APIPlaintextFrameHelper(APIFrameHelper):
             if (msg_type := self._read_varuint()) == -1:
                 return
 
+            packet_data: bytes | None
             if length == 0:
-                packet_data = b""
+                packet_data = EMPTY_PACKET
             else:
                 # The packet data is not yet available, wait for more data
                 # to arrive before continuing, since callback_packet has not
                 # been called yet the buffer will not be cleared and the next
                 # call to data_received will continue processing the packet
                 # at the start of the frame.
-                if (maybe_packet_data := self._read(length)) is None:
+                if (packet_data := self._read(length)) is None:
                     return
-                packet_data = maybe_packet_data
+                if TYPE_CHECKING:
+                    assert packet_data is not None, "Packet data should be set"
 
             self._remove_from_buffer()
             self._connection.process_packet(msg_type, packet_data)