From cea62efbbae4bb688561579c061b32b7dea9f81c Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 26 Nov 2023 13:25:55 -0600 Subject: [PATCH] Small cython cleanps to the frame helpers (#733) --- aioesphomeapi/_frame_helper/base.pxd | 4 ++-- aioesphomeapi/_frame_helper/base.py | 3 ++- aioesphomeapi/_frame_helper/noise.pxd | 4 ++-- aioesphomeapi/_frame_helper/noise.py | 4 ++-- aioesphomeapi/_frame_helper/plain_text.pxd | 6 ++++-- aioesphomeapi/_frame_helper/plain_text.py | 11 ++++++++--- 6 files changed, 20 insertions(+), 12 deletions(-) diff --git a/aioesphomeapi/_frame_helper/base.pxd b/aioesphomeapi/_frame_helper/base.pxd index e70acab..2f16b11 100644 --- a/aioesphomeapi/_frame_helper/base.pxd +++ b/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) diff --git a/aioesphomeapi/_frame_helper/base.py b/aioesphomeapi/_frame_helper/base.py index 8ed4392..597f41b 100644 --- a/aioesphomeapi/_frame_helper/base.py +++ b/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()) diff --git a/aioesphomeapi/_frame_helper/noise.pxd b/aioesphomeapi/_frame_helper/noise.pxd index 1604de7..c264eb3 100644 --- a/aioesphomeapi/_frame_helper/noise.pxd +++ b/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) diff --git a/aioesphomeapi/_frame_helper/noise.py b/aioesphomeapi/_frame_helper/noise.py index 8622ff7..41ea06f 100644 --- a/aioesphomeapi/_frame_helper/noise.py +++ b/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 diff --git a/aioesphomeapi/_frame_helper/plain_text.pxd b/aioesphomeapi/_frame_helper/plain_text.pxd index 382b3b8..ac1a93d 100644 --- a/aioesphomeapi/_frame_helper/plain_text.pxd +++ b/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) diff --git a/aioesphomeapi/_frame_helper/plain_text.py b/aioesphomeapi/_frame_helper/plain_text.py index 7a4beee..7be442b 100644 --- a/aioesphomeapi/_frame_helper/plain_text.py +++ b/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)