From 94157a8ac101f476fec58222395dd1b2c3570e4e Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 9 Nov 2023 10:20:42 -0600
Subject: [PATCH] Small speed up to the plaintext frame helper (#625)

---
 aioesphomeapi/_frame_helper/noise.py      |  3 +--
 aioesphomeapi/_frame_helper/plain_text.py | 25 +++++++++++++----------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/aioesphomeapi/_frame_helper/noise.py b/aioesphomeapi/_frame_helper/noise.py
index b0e4d34..bf356fc 100644
--- a/aioesphomeapi/_frame_helper/noise.py
+++ b/aioesphomeapi/_frame_helper/noise.py
@@ -142,8 +142,7 @@ class APINoiseFrameHelper(APIFrameHelper):
         self._add_to_buffer(data)
         while self._buffer:
             self._pos = 0
-            header = self._read_exactly(3)
-            if header is None:
+            if (header := self._read_exactly(3)) is None:
                 return
             preamble = header[0]
             msg_size_high = header[1]
diff --git a/aioesphomeapi/_frame_helper/plain_text.py b/aioesphomeapi/_frame_helper/plain_text.py
index b629b3e..df1a4f4 100644
--- a/aioesphomeapi/_frame_helper/plain_text.py
+++ b/aioesphomeapi/_frame_helper/plain_text.py
@@ -78,15 +78,16 @@ class APIPlaintextFrameHelper(APIFrameHelper):
                 f"{self._log_name}: Error while writing data: {err}"
             ) from err
 
-    def data_received(self, data: bytes) -> None:  # pylint: disable=too-many-branches
+    def data_received(  # pylint: disable=too-many-branches,too-many-return-statements
+        self, data: bytes
+    ) -> None:
         self._add_to_buffer(data)
         while self._buffer:
             # Read preamble, which should always 0x00
             # Also try to get the length and msg type
             # to avoid multiple calls to _read_exactly
             self._pos = 0
-            init_bytes = self._read_exactly(3)
-            if init_bytes is None:
+            if (init_bytes := self._read_exactly(3)) is None:
                 return
             msg_type_int: int | None = None
             length_int = 0
@@ -125,22 +126,25 @@ class APIPlaintextFrameHelper(APIFrameHelper):
                 length = init_bytes[1:3]
                 # If the message is long, we need to read the rest of the length
                 while length[-1] & 0x80 == 0x80:
-                    add_length = self._read_exactly(1)
-                    if add_length is None:
+                    if (add_length := self._read_exactly(1)) is None:
                         return
                     length += add_length
                 length_int = bytes_to_varuint(length) or 0
                 # Since the length is longer than 1 byte we do not have the
                 # message type yet.
-                msg_type = b""
+                if (msg_type_byte := self._read_exactly(1)) is None:
+                    return
+                msg_type = msg_type_byte
+                if msg_type[-1] & 0x80 != 0x80:
+                    # Message type is only 1 byte
+                    msg_type_int = msg_type[0]
 
             # If the we do not have the message type yet because the message
             # length was so long it did not fit into the first byte we need
             # to read the (rest) of the message type
             if msg_type_int is None:
-                while not msg_type or msg_type[-1] & 0x80 == 0x80:
-                    add_msg_type = self._read_exactly(1)
-                    if add_msg_type is None:
+                while msg_type[-1] & 0x80 == 0x80:
+                    if (add_msg_type := self._read_exactly(1)) is None:
                         return
                     msg_type += add_msg_type
                 msg_type_int = bytes_to_varuint(msg_type)
@@ -151,13 +155,12 @@ class APIPlaintextFrameHelper(APIFrameHelper):
             if length_int == 0:
                 packet_data = b""
             else:
-                maybe_packet_data = self._read_exactly(length_int)
                 # 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 is None:
+                if (maybe_packet_data := self._read_exactly(length_int)) is None:
                     return
                 packet_data = maybe_packet_data