diff --git a/aioesphomeapi/client.py b/aioesphomeapi/client.py index e337d29..046766a 100644 --- a/aioesphomeapi/client.py +++ b/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) @@ -720,13 +722,9 @@ class APIClient: timeout: float, ) -> message.Message: """Send a BluetoothDeviceRequest and wait for a response.""" + req = BluetoothDeviceRequest(address=address, request_type=request_type) [response] = await self._get_connection().send_messages_await_response_complex( - ( - BluetoothDeviceRequest( - address=address, - request_type=request_type, - ), - ), + (req,), predicate_func, predicate_func, msg_types, @@ -737,42 +735,18 @@ class APIClient: async def bluetooth_gatt_get_services( self, address: int ) -> ESPHomeBluetoothGATTServices: - append_types = ( - BluetoothDeviceConnectionResponse, - BluetoothGATTGetServicesResponse, - BluetoothGATTErrorResponse, - ) - stop_types = ( - BluetoothDeviceConnectionResponse, - BluetoothGATTGetServicesDoneResponse, - BluetoothGATTErrorResponse, - ) + error_types = (BluetoothGATTErrorResponse, BluetoothDeviceConnectionResponse) + append_types = (*error_types, BluetoothGATTGetServicesResponse) + stop_types = (*error_types, BluetoothGATTGetServicesDoneResponse) msg_types = ( BluetoothGATTGetServicesResponse, 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, ) diff --git a/aioesphomeapi/client_callbacks.py b/aioesphomeapi/client_callbacks.py index 4b378ae..f4cc3f2 100644 --- a/aioesphomeapi/client_callbacks.py +++ b/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,23 @@ 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 + | BluetoothGATTGetServicesResponse + | BluetoothGATTGetServicesDoneResponse + | BluetoothGATTErrorResponse, +) -> bool: + """Filter Bluetooth messages of a specific type and address.""" + return type(msg) in msg_types and bool(msg.address == address)