From bc1b7651049329fccdf1e2bd8c19fb80c0961ae2 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 5 May 2024 21:19:03 -0500
Subject: [PATCH] Cache bytes to protobuf

---
 aioesphomeapi/connection.py | 38 +++++++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py
index adc2101..0e995c7 100644
--- a/aioesphomeapi/connection.py
+++ b/aioesphomeapi/connection.py
@@ -180,6 +180,19 @@ def handle_complex_message(
 _handle_complex_message = handle_complex_message
 
 
+@lru_cache(maxsize=1024)
+def _message_from_data(msg_type_proto: _int, data: _bytes) -> message.Message | None:
+    """Create a message from data."""
+    if (klass := MESSAGE_TYPE_TO_PROTO.get(msg_type_proto)) is None:
+        return None
+    msg: message.Message = klass()
+    # MergeFromString instead of ParseFromString since
+    # ParseFromString will clear the message first and
+    # the msg is already empty.
+    msg.MergeFromString(data)
+    return msg
+
+
 class APIConnection:
     """This class represents _one_ connection to a remote native API device.
 
@@ -876,22 +889,10 @@ class APIConnection:
     def process_packet(self, msg_type_proto: _int, data: _bytes) -> None:
         """Process an incoming packet."""
         debug_enabled = self._debug_enabled
-        if (klass := MESSAGE_TYPE_TO_PROTO.get(msg_type_proto)) is None:
-            if debug_enabled:
-                _LOGGER.debug(
-                    "%s: Skipping unknown message type %s",
-                    self.log_name,
-                    msg_type_proto,
-                )
-            return
-
         try:
-            msg: message.Message = klass()
-            # MergeFromString instead of ParseFromString since
-            # ParseFromString will clear the message first and
-            # the msg is already empty.
-            msg.MergeFromString(data)
+            msg = _message_from_data(msg_type_proto, data)
         except Exception as e:
+            klass = MESSAGE_TYPE_TO_PROTO[msg_type_proto]
             _LOGGER.error(
                 "%s: Invalid protobuf message: type=%s data=%s: %s",
                 self.log_name,
@@ -907,6 +908,15 @@ class APIConnection:
             )
             raise
 
+        if msg is None:
+            if debug_enabled:
+                _LOGGER.debug(
+                    "%s: Skipping unknown message type %s",
+                    self.log_name,
+                    msg_type_proto,
+                )
+            return
+
         msg_type = type(msg)
 
         if debug_enabled: