From 83b4f4361082f758af22de40bcdac62f7ece8aa5 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 23 Nov 2023 19:20:52 +0100 Subject: [PATCH] Refactor frame helper to get debug state from connection (#679) --- aioesphomeapi/_frame_helper/base.pxd | 5 ++--- aioesphomeapi/_frame_helper/base.py | 11 +++++------ aioesphomeapi/_frame_helper/noise.pxd | 2 +- aioesphomeapi/_frame_helper/noise.py | 11 +++++++---- aioesphomeapi/_frame_helper/plain_text.pxd | 2 +- aioesphomeapi/_frame_helper/plain_text.py | 6 ++++-- aioesphomeapi/connection.py | 2 +- bench/raw_ble_plain_text_with_callback.py | 1 + tests/test__frame_helper.py | 2 +- 9 files changed, 23 insertions(+), 19 deletions(-) diff --git a/aioesphomeapi/_frame_helper/base.pxd b/aioesphomeapi/_frame_helper/base.pxd index babfbb4..db554d5 100644 --- a/aioesphomeapi/_frame_helper/base.pxd +++ b/aioesphomeapi/_frame_helper/base.pxd @@ -19,7 +19,6 @@ cdef class APIFrameHelper: cdef unsigned int _pos cdef object _client_info cdef str _log_name - cdef object _debug_enabled cpdef set_log_name(self, str log_name) @@ -32,6 +31,6 @@ cdef class APIFrameHelper: @cython.locals(end_of_frame_pos="unsigned int") cdef _remove_from_buffer(self) - cpdef write_packets(self, list packets) + cpdef write_packets(self, list packets, bint debug_enabled) - cdef _write_bytes(self, bytes data) + cdef _write_bytes(self, bytes data, bint debug_enabled) diff --git a/aioesphomeapi/_frame_helper/base.py b/aioesphomeapi/_frame_helper/base.py index 55f1a14..2ad4b36 100644 --- a/aioesphomeapi/_frame_helper/base.py +++ b/aioesphomeapi/_frame_helper/base.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import logging from abc import abstractmethod -from functools import partial from typing import TYPE_CHECKING, Callable, cast from ..core import HandshakeAPIError, SocketClosedAPIError @@ -39,7 +38,6 @@ class APIFrameHelper: "_pos", "_client_info", "_log_name", - "_debug_enabled", ) def __init__( @@ -60,7 +58,6 @@ class APIFrameHelper: self._pos = 0 self._client_info = client_info self._log_name = log_name - self._debug_enabled = partial(_LOGGER.isEnabledFor, logging.DEBUG) def set_log_name(self, log_name: str) -> None: """Set the log name.""" @@ -139,7 +136,9 @@ class APIFrameHelper: handshake_handle.cancel() @abstractmethod - def write_packets(self, packets: list[tuple[int, bytes]]) -> None: + def write_packets( + self, packets: list[tuple[int, bytes]], debug_enabled: bool + ) -> None: """Write a packets to the socket. Packets are in the format of tuple[protobuf_type, protobuf_data] @@ -181,9 +180,9 @@ class APIFrameHelper: def resume_writing(self) -> None: """Stub.""" - def _write_bytes(self, data: bytes) -> None: + def _write_bytes(self, data: bytes, debug_enabled: bool) -> None: """Write bytes to the socket.""" - if self._debug_enabled() is True: + if debug_enabled: _LOGGER.debug("%s: Sending frame: [%s]", self._log_name, data.hex()) if TYPE_CHECKING: diff --git a/aioesphomeapi/_frame_helper/noise.pxd b/aioesphomeapi/_frame_helper/noise.pxd index fe32420..41923e0 100644 --- a/aioesphomeapi/_frame_helper/noise.pxd +++ b/aioesphomeapi/_frame_helper/noise.pxd @@ -51,6 +51,6 @@ cdef class APINoiseFrameHelper(APIFrameHelper): frame_len=cython.uint, type_=object ) - cpdef write_packets(self, list packets) + cpdef 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 b85c90f..1416df5 100644 --- a/aioesphomeapi/_frame_helper/noise.py +++ b/aioesphomeapi/_frame_helper/noise.py @@ -1,6 +1,7 @@ from __future__ import annotations import binascii +import logging from functools import partial from struct import Struct from typing import TYPE_CHECKING, Any, Callable @@ -20,7 +21,7 @@ from ..core import ( InvalidEncryptionKeyAPIError, ProtocolAPIError, ) -from .base import APIFrameHelper +from .base import _LOGGER, APIFrameHelper if TYPE_CHECKING: from ..connection import APIConnection @@ -180,7 +181,7 @@ class APINoiseFrameHelper(APIFrameHelper): frame_len = len(handshake_frame) header = bytes((0x01, (frame_len >> 8) & 0xFF, frame_len & 0xFF)) hello_handshake = NOISE_HELLO + header + handshake_frame - self._write_bytes(hello_handshake) + self._write_bytes(hello_handshake, _LOGGER.isEnabledFor(logging.DEBUG)) def _handle_hello(self, server_hello: bytes) -> None: """Perform the handshake with the server.""" @@ -284,7 +285,9 @@ class APINoiseFrameHelper(APIFrameHelper): ) self._ready_future.set_result(None) - def write_packets(self, packets: list[tuple[int, bytes]]) -> None: + def write_packets( + self, packets: list[tuple[int, bytes]], debug_enabled: bool + ) -> None: """Write a packets to the socket. Packets are in the format of tuple[protobuf_type, protobuf_data] @@ -314,7 +317,7 @@ class APINoiseFrameHelper(APIFrameHelper): out.append(header) out.append(frame) - self._write_bytes(b"".join(out)) + self._write_bytes(b"".join(out), debug_enabled) def _handle_frame(self, frame: bytes) -> None: """Handle an incoming frame.""" diff --git a/aioesphomeapi/_frame_helper/plain_text.pxd b/aioesphomeapi/_frame_helper/plain_text.pxd index 4591a8b..6e5fe92 100644 --- a/aioesphomeapi/_frame_helper/plain_text.pxd +++ b/aioesphomeapi/_frame_helper/plain_text.pxd @@ -35,4 +35,4 @@ cdef class APIPlaintextFrameHelper(APIFrameHelper): packet=tuple, type_=object ) - cpdef write_packets(self, list packets) + cpdef 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 82a5b51..19618f7 100644 --- a/aioesphomeapi/_frame_helper/plain_text.py +++ b/aioesphomeapi/_frame_helper/plain_text.py @@ -56,7 +56,9 @@ class APIPlaintextFrameHelper(APIFrameHelper): super().connection_made(transport) self._ready_future.set_result(None) - def write_packets(self, packets: list[tuple[int, bytes]]) -> None: + def write_packets( + self, packets: list[tuple[int, bytes]], debug_enabled: bool + ) -> None: """Write a packets to the socket. Packets are in the format of tuple[protobuf_type, protobuf_data] @@ -72,7 +74,7 @@ class APIPlaintextFrameHelper(APIFrameHelper): out.append(varuint_to_bytes(type_)) out.append(data) - self._write_bytes(b"".join(out)) + self._write_bytes(b"".join(out), debug_enabled) def data_received( # pylint: disable=too-many-branches,too-many-return-statements self, data: bytes | bytearray | memoryview diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py index 4d97065..d81738f 100644 --- a/aioesphomeapi/connection.py +++ b/aioesphomeapi/connection.py @@ -643,7 +643,7 @@ class APIConnection: assert self._frame_helper is not None try: - self._frame_helper.write_packets(packets) + self._frame_helper.write_packets(packets, debug_enabled) except SocketAPIError as err: # If writing packet fails, we don't know what state the frames # are in anymore and we have to close the connection diff --git a/bench/raw_ble_plain_text_with_callback.py b/bench/raw_ble_plain_text_with_callback.py index dda9713..a8debb8 100644 --- a/bench/raw_ble_plain_text_with_callback.py +++ b/bench/raw_ble_plain_text_with_callback.py @@ -8,6 +8,7 @@ from aioesphomeapi.api_pb2 import ( ) from aioesphomeapi.client import APIClient from aioesphomeapi.client_callbacks import on_ble_raw_advertisement_response + # cythonize -X language_level=3 -a -i aioesphomeapi/client_callbacks.py # cythonize -X language_level=3 -a -i aioesphomeapi/connection.py diff --git a/tests/test__frame_helper.py b/tests/test__frame_helper.py index e68df21..4956b54 100644 --- a/tests/test__frame_helper.py +++ b/tests/test__frame_helper.py @@ -541,7 +541,7 @@ async def test_noise_frame_helper_handshake_success_with_single_packet(): assert not writes await handshake_task - helper.write_packets([(1, b"to device")]) + helper.write_packets([(1, b"to device")], True) encrypted_packet = writes.pop() header = encrypted_packet[0:1] assert header == b"\x01"