Reduce duplicate Bluetooth message filtering code

aioesphomeapi/client.py

@@ -76,8 +76,9 @@ from .client_callbacks import (
     on_bluetooth_connections_free_response,
     on_bluetooth_device_connection_response,
     on_bluetooth_gatt_notify_data_response,
+    on_bluetooth_handle_message,
     on_bluetooth_le_advertising_response,
-    on_bluetooth_message,
+    on_bluetooth_message_types,
     on_home_assistant_service_response,
     on_state_msg,
     on_subscribe_home_assistant_state_response,
@@ -465,7 +466,7 @@ class APIClient:
         ),
         timeout: float = 10.0,
     ) -> message.Message:
-        message_filter = partial(on_bluetooth_message, address, handle)
+        message_filter = partial(on_bluetooth_handle_message, address, handle)
         msg_types = (response_type, BluetoothGATTErrorResponse)
         [resp] = await self._get_connection().send_messages_await_response_complex(
             (request,),
@@ -684,11 +685,12 @@ class APIClient:
         timeout: float,
     ) -> message.Message:
         """Send a BluetoothDeviceRequest watch for the connection state to change."""
+        types_with_response = (BluetoothDeviceConnectionResponse, *msg_types)
         response = await self._bluetooth_device_request(
             address,
             request_type,
-            lambda msg: msg.address == address,
-            (BluetoothDeviceConnectionResponse, *msg_types),
+            partial(on_bluetooth_message_types, address, types_with_response),
+            types_with_response,
             timeout,
         )
         self._raise_for_ble_connection_change(address, response, msg_types)
@@ -752,27 +754,10 @@ class APIClient:
             BluetoothGATTGetServicesDoneResponse,
             BluetoothGATTErrorResponse,
         )
-
-        def do_append(
-            msg: BluetoothDeviceConnectionResponse
-            | BluetoothGATTGetServicesResponse
-            | BluetoothGATTGetServicesDoneResponse
-            | BluetoothGATTErrorResponse,
-        ) -> bool:
-            return type(msg) in append_types and msg.address == address
-
-        def do_stop(
-            msg: BluetoothDeviceConnectionResponse
-            | BluetoothGATTGetServicesResponse
-            | BluetoothGATTGetServicesDoneResponse
-            | BluetoothGATTErrorResponse,
-        ) -> bool:
-            return type(msg) in stop_types and msg.address == address
-
         resp = await self._get_connection().send_messages_await_response_complex(
             (BluetoothGATTGetServicesRequest(address=address),),
-            do_append,
-            do_stop,
+            partial(on_bluetooth_message_types, address, append_types),
+            partial(on_bluetooth_message_types, address, stop_types),
             (*msg_types, BluetoothDeviceConnectionResponse),
             DEFAULT_BLE_TIMEOUT,
         )

aioesphomeapi/client_callbacks.py

@@ -10,6 +10,8 @@ from .api_pb2 import (  # type: ignore
     BluetoothConnectionsFreeResponse,
     BluetoothDeviceConnectionResponse,
     BluetoothGATTErrorResponse,
+    BluetoothGATTGetServicesDoneResponse,
+    BluetoothGATTGetServicesResponse,
     BluetoothGATTNotifyDataResponse,
     BluetoothGATTNotifyResponse,
     BluetoothGATTReadResponse,
@@ -118,7 +120,7 @@ def on_bluetooth_device_connection_response(
             connect_future.set_result(None)
 
 
-def on_bluetooth_message(
+def on_bluetooth_handle_message(
     address: int,
     handle: int,
     msg: BluetoothGATTErrorResponse
@@ -127,7 +129,24 @@ def on_bluetooth_message(
     | BluetoothGATTWriteResponse
     | BluetoothDeviceConnectionResponse,
 ) -> bool:
-    """Handle a Bluetooth message."""
+    """Filter a Bluetooth message for an address and handle."""
     if type(msg) is BluetoothDeviceConnectionResponse:
         return bool(msg.address == address)
     return bool(msg.address == address and msg.handle == handle)
+
+
+def on_bluetooth_message_types(
+    address: int,
+    msg_types: tuple[type[message.Message]],
+    msg: BluetoothGATTErrorResponse
+    | BluetoothGATTNotifyResponse
+    | BluetoothGATTReadResponse
+    | BluetoothGATTWriteResponse
+    | BluetoothDeviceConnectionResponse
+    | BluetoothDeviceConnectionResponse
+    | BluetoothGATTGetServicesResponse
+    | BluetoothGATTGetServicesDoneResponse
+    | BluetoothGATTErrorResponse,
+) -> bool:
+    """Filter Bluetooth messages of a specific type and address."""
+    return type(msg) in msg_types and bool(msg.address == address)