Guard client connection asserts with if TYPE_CHECKING (#659)

aioesphomeapi/client.py

@@ -388,6 +388,7 @@ class APIClient:
         login: bool = False,
     ) -> None:
         """Finish connecting to the device."""
+        if TYPE_CHECKING:
             assert self._connection is not None
         try:
             await self._connection.finish_connection(login=login)
@@ -423,6 +424,7 @@ class APIClient:
     async def device_info(self) -> DeviceInfo:
         self._check_authenticated()
         connection = self._connection
+        if TYPE_CHECKING:
             assert connection is not None
         resp = await connection.send_message_await_response(
             DeviceInfoRequest(), DeviceInfoResponse
@@ -449,6 +451,7 @@ class APIClient:
         def do_stop(msg: message.Message) -> bool:
             return isinstance(msg, ListEntitiesDoneResponse)
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         resp = await self._connection.send_messages_await_response_complex(
             (ListEntitiesRequest(),), do_append, do_stop, msg_types, 60
@@ -491,6 +494,7 @@ class APIClient:
                     del image_stream[msg_key]
                     on_state(CameraState(key=msg.key, data=image_data))  # type: ignore[call-arg]
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message_callback_response(
             SubscribeStatesRequest(), _on_state_msg, msg_types
@@ -508,6 +512,7 @@ class APIClient:
             req.level = log_level
         if dump_config is not None:
             req.dump_config = dump_config
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message_callback_response(
             req, on_log, (SubscribeLogsResponse,)
@@ -523,6 +528,7 @@ class APIClient:
         ) -> None:
             on_service_call(HomeassistantServiceCall.from_pb(msg))
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message_callback_response(
             SubscribeHomeassistantServicesRequest(),
@@ -563,6 +569,7 @@ class APIClient:
     ) -> message.Message:
         self._check_authenticated()
         msg_types = (response_type, BluetoothGATTErrorResponse)
+        if TYPE_CHECKING:
             assert self._connection is not None
 
         message_filter = partial(self._filter_bluetooth_message, address, handle)
@@ -586,6 +593,7 @@ class APIClient:
         ) -> None:
             on_bluetooth_le_advertisement(BluetoothLEAdvertisement.from_pb(msg))  # type: ignore[misc]
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         unsub_callback = self._connection.send_message_callback_response(
             SubscribeBluetoothLEAdvertisementsRequest(flags=0),
@@ -608,6 +616,7 @@ class APIClient:
         self._check_authenticated()
         msg_types = (BluetoothLERawAdvertisementsResponse,)
 
+        if TYPE_CHECKING:
             assert self._connection is not None
 
         def _on_ble_raw_advertisement_response(
@@ -643,6 +652,7 @@ class APIClient:
         ) -> None:
             on_bluetooth_connections_free_update(msg.free, msg.limit)
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         return self._connection.send_message_callback_response(
             SubscribeBluetoothConnectionsFreeRequest(),
@@ -687,7 +697,9 @@ class APIClient:
         debug = _LOGGER.isEnabledFor(logging.DEBUG)
         connect_future: asyncio.Future[None] = self._loop.create_future()
 
+        if TYPE_CHECKING:
             assert self._connection is not None
+
         if has_cache:
             # REMOTE_CACHING feature with cache: requestor has services and mtu cached
             request_type = BluetoothDeviceRequestType.CONNECT_V3_WITH_CACHE
@@ -854,6 +866,7 @@ class APIClient:
         timeout: float,
     ) -> message.Message:
         self._check_authenticated()
+        if TYPE_CHECKING:
             assert self._connection is not None
         [response] = await self._connection.send_messages_await_response_complex(
             (
@@ -887,6 +900,7 @@ class APIClient:
         def do_stop(msg: message.Message) -> bool:
             return isinstance(msg, stop_types) and msg.address == address
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         resp = await self._connection.send_messages_await_response_complex(
             (BluetoothGATTGetServicesRequest(address=address),),
@@ -932,6 +946,7 @@ class APIClient:
         req.data = data
 
         if not response:
+            if TYPE_CHECKING:
                 assert self._connection is not None
             self._connection.send_message(req)
             return
@@ -993,6 +1008,7 @@ class APIClient:
         req.data = data
 
         if not wait_for_response:
+            if TYPE_CHECKING:
                 assert self._connection is not None
             self._connection.send_message(req)
             return
@@ -1035,6 +1051,7 @@ class APIClient:
             if address == msg.address and handle == msg.handle:
                 on_bluetooth_gatt_notify(handle, bytearray(msg.data))
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         remove_callback = self._connection.add_message_callback(
             _on_bluetooth_gatt_notify_data_response, (BluetoothGATTNotifyDataResponse,)
@@ -1064,6 +1081,7 @@ class APIClient:
         ) -> None:
             on_state_sub(msg.entity_id, msg.attribute)
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message_callback_response(
             SubscribeHomeAssistantStatesRequest(),
@@ -1076,6 +1094,7 @@ class APIClient:
     ) -> None:
         self._check_authenticated()
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(
             HomeAssistantStateResponse(
@@ -1116,6 +1135,7 @@ class APIClient:
             elif position == 0.0:
                 req.legacy_command = LegacyCoverCommand.CLOSE
                 req.has_legacy_command = True
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1147,10 +1167,11 @@ class APIClient:
         if direction is not None:
             req.has_direction = True
             req.direction = direction
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
-    async def light_command(
+    async def light_command(  # pylint: disable=too-many-branches
         self,
         key: int,
         state: bool | None = None,
@@ -1208,6 +1229,7 @@ class APIClient:
         if effect is not None:
             req.has_effect = True
             req.effect = effect
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1217,6 +1239,7 @@ class APIClient:
         req = SwitchCommandRequest()
         req.key = key
         req.state = state
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1269,6 +1292,7 @@ class APIClient:
         if custom_preset is not None:
             req.has_custom_preset = True
             req.custom_preset = custom_preset
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1278,6 +1302,7 @@ class APIClient:
         req = NumberCommandRequest()
         req.key = key
         req.state = state
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1287,6 +1312,7 @@ class APIClient:
         req = SelectCommandRequest()
         req.key = key
         req.state = state
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1314,6 +1340,7 @@ class APIClient:
         if duration is not None:
             req.duration = duration
             req.has_duration = True
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1322,6 +1349,7 @@ class APIClient:
 
         req = ButtonCommandRequest()
         req.key = key
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1338,6 +1366,7 @@ class APIClient:
         req.command = command
         if code is not None:
             req.code = code
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1362,6 +1391,7 @@ class APIClient:
         if media_url is not None:
             req.media_url = media_url
             req.has_media_url = True
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1371,6 +1401,7 @@ class APIClient:
         req = TextCommandRequest()
         req.key = key
         req.state = state
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1409,6 +1440,7 @@ class APIClient:
             args.append(arg)
         # pylint: disable=no-member
         req.args.extend(args)
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1418,6 +1450,7 @@ class APIClient:
         req = CameraImageRequest()
         req.single = single
         req.stream = stream
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1478,6 +1511,7 @@ class APIClient:
                 self._background_tasks.add(stop_task)
                 stop_task.add_done_callback(self._background_tasks.discard)
 
+        if TYPE_CHECKING:
             assert self._connection is not None
 
         self._connection.send_message(SubscribeVoiceAssistantRequest(subscribe=True))
@@ -1517,6 +1551,7 @@ class APIClient:
         # pylint: disable=no-member
         req.data.extend(data_args)
 
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)
 
@@ -1533,5 +1568,6 @@ class APIClient:
         req.command = command
         if code is not None:
             req.code = code
+        if TYPE_CHECKING:
             assert self._connection is not None
         self._connection.send_message(req)