pymammotion 0.2.62__py3-none-any.whl → 0.5.51__py3-none-any.whl

pymammotion/mammotion/devices/mammotion_bluetooth.py

@@ -1,9 +1,11 @@
 import asyncio
+from collections.abc import Awaitable, Callable
 import logging
-from typing import Any, cast
+import time
+from typing import Any
 from uuid import UUID
 
-import betterproto
+import betterproto2
 from bleak import BleakGATTCharacteristic, BleakGATTServiceCollection, BLEDevice
 from bleak.exc import BleakDBusError
 from bleak_retry_connector import (
@@ -13,12 +15,12 @@ from bleak_retry_connector import (
     establish_connection,
 )
 
+from pymammotion.aliyun.model.dev_by_account_response import Device
 from pymammotion.bluetooth import BleMessage
-from pymammotion.data.model.device import MowingDevice
+from pymammotion.data.mower_state_manager import MowerStateManager
 from pymammotion.mammotion.commands.mammotion_command import MammotionCommand
 from pymammotion.mammotion.devices.base import MammotionBaseDevice
-from pymammotion.proto import has_field
-from pymammotion.proto.luba_msg import LubaMsg
+from pymammotion.proto import LubaMsg
 
 DBUS_ERROR_BACKOFF_TIME = 0.25
 
@@ -69,33 +71,61 @@ async def _handle_retry(fut: asyncio.Future[None], func, command: bytes) -> None
 class MammotionBaseBLEDevice(MammotionBaseDevice):
     """Base class for Mammotion BLE devices."""
 
-    def __init__(self, mowing_state: MowingDevice, device: BLEDevice, interface: int = 0, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        state_manager: MowerStateManager,
+        cloud_device: Device,
+        device: BLEDevice,
+        interface: int = 0,
+        **kwargs: Any,
+    ) -> None:
         """Initialize MammotionBaseBLEDevice."""
-        super().__init__(mowing_state)
+        super().__init__(state_manager, cloud_device)
+        self.command_sent_time = 0
         self._disconnect_strategy = True
         self._ble_sync_task = None
         self._prev_notification = None
         self._interface = f"hci{interface}"
-        self._device = device
+        self.ble_device = device
         self._client: BleakClientWithServiceCache | None = None
         self._read_char: BleakGATTCharacteristic | int | str | UUID = 0
         self._write_char: BleakGATTCharacteristic | int | str | UUID = 0
         self._disconnect_timer: asyncio.TimerHandle | None = None
         self._message: BleMessage | None = None
-        self._commands: MammotionCommand = MammotionCommand(device.name)
+        self._commands: MammotionCommand = MammotionCommand(device.name, 1)
+        self.command_queue = asyncio.Queue()
         self._expected_disconnect = False
         self._connect_lock = asyncio.Lock()
         self._operation_lock = asyncio.Lock()
         self._key: str | None = None
+        self._cloud_device = cloud_device
         self.set_queue_callback(self.queue_command)
+        loop = asyncio.get_event_loop()
+        loop.create_task(self.process_queue())
+
+    def __del__(self) -> None:
+        """Cleanup."""
+        if self._disconnect_timer:
+            self._disconnect_timer.cancel()
+        if self._ble_sync_task:
+            self._ble_sync_task.cancel()
+
+        self._state_manager.ble_queue_command_callback.remove_subscribers(self.queue_command)
+
+    def set_notification_callback(self, func: Callable[[tuple[str, Any | None]], Awaitable[None]]) -> None:
+        self._state_manager.ble_on_notification_callback.add_subscribers(func)
+
+    def set_queue_callback(self, func: Callable[[str, dict[str, Any]], Awaitable[None]]) -> None:
+        self._state_manager.ble_queue_command_callback.add_subscribers(func)
 
     def update_device(self, device: BLEDevice) -> None:
         """Update the BLE device."""
-        self._device = device
+        self.ble_device = device
 
     async def _ble_sync(self) -> None:
-        command_bytes = self._commands.send_todev_ble_sync(2)
-        await self._message.post_custom_data_bytes(command_bytes)
+        if self._client is not None and self._client.is_connected:
+            command_bytes = self._commands.send_todev_ble_sync(2)
+            await self._message.post_custom_data_bytes(command_bytes)
 
     async def run_periodic_sync_task(self) -> None:
         """Send ble sync to robot."""
@@ -113,11 +143,37 @@ class MammotionBaseBLEDevice(MammotionBaseDevice):
 
     async def stop(self) -> None:
         """Stop all tasks and disconnect."""
-        self._ble_sync_task.cancel()
-        await self._client.disconnect()
-
-    async def queue_command(self, key: str, **kwargs: Any) -> bytes | None:
-        return await self._send_command_with_args(key, **kwargs)
+        if self._ble_sync_task:
+            self._ble_sync_task.cancel()
+        if self._client is not None and self._client.is_connected:
+            await self._client.disconnect()
+
+    async def queue_command(self, key: str, **kwargs: Any) -> None:
+        # Create a future to hold the result
+        _LOGGER.debug("Queueing command: %s", key)
+        future = asyncio.Future()
+        # Put the command in the queue as a tuple (key, command, future)
+        command_bytes = getattr(self._commands, key)(**kwargs)
+        await self.command_queue.put((key, command_bytes, future))
+        # Wait for the future to be resolved
+        await future
+        # return await self._send_command_with_args(key, **kwargs)
+
+    async def process_queue(self) -> None:
+        while True:
+            # Get the next item from the queue
+            key, command, future = await self.command_queue.get()
+            try:
+                # Process the command using _execute_command_locked
+                result = await self._send_command_locked(key, command)
+                # Set the result on the future
+                future.set_result(result)
+            except Exception as ex:
+                # Set the exception on the future if something goes wrong
+                future.set_exception(ex)
+            finally:
+                # Mark the task as done
+                self.command_queue.task_done()
 
     async def _send_command_with_args(self, key: str, **kwargs) -> bytes | None:
         """Send command to device and read response."""
@@ -184,15 +240,17 @@ class MammotionBaseBLEDevice(MammotionBaseDevice):
     @property
     def name(self) -> str:
         """Return device name."""
-        return f"{self._device.name} ({self._device.address})"
+        return f"{self.ble_device.name} ({self.ble_device.address})"
 
     @property
     def rssi(self) -> int:
         """Return RSSI of device."""
-        try:
-            return cast(self.mower.sys.toapp_report_data.connect.ble_rssi, int)
-        finally:
-            return 0
+        return self.mower.report_data.connect.ble_rssi
+
+    @property
+    def client(self) -> BleakClientWithServiceCache | None:
+        """Return client."""
+        return self._client
 
     async def _ensure_connected(self) -> None:
         """Ensure connection to device is established."""
@@ -223,11 +281,11 @@ class MammotionBaseBLEDevice(MammotionBaseDevice):
             _LOGGER.debug("%s: Connecting; RSSI: %s", self.name, self.rssi)
             client: BleakClientWithServiceCache = await establish_connection(
                 BleakClientWithServiceCache,
-                self._device,
+                self.ble_device,
                 self.name,
                 self._disconnected,
                 max_attempts=10,
-                ble_device_callback=lambda: self._device,
+                ble_device_callback=lambda: self.ble_device,
             )
             _LOGGER.debug("%s: Connected; RSSI: %s", self.name, self.rssi)
             self._client = client
@@ -255,11 +313,10 @@ class MammotionBaseBLEDevice(MammotionBaseDevice):
             )
             self._reset_disconnect_timer()
             await self._start_notify()
-            command_bytes = self._commands.send_todev_ble_sync(2)
-            await self._message.post_custom_data_bytes(command_bytes)
+            await self._ble_sync()
             self.schedule_ble_sync()
 
-    async def _send_command_locked(self, key: str, command: bytes) -> bytes:
+    async def _send_command_locked(self, key: str, command: bytes) -> None:
         """Send command to device and read response."""
         await self._ensure_connected()
         try:
@@ -284,73 +341,60 @@ class MammotionBaseBLEDevice(MammotionBaseDevice):
 
     async def _notification_handler(self, _sender: BleakGATTCharacteristic, data: bytearray) -> None:
         """Handle notification responses."""
+
         if self._message is None:
             return
         result = self._message.parseNotification(data)
         if result == 0:
             data = await self._message.parseBlufiNotifyData(True)
-            self._update_raw_data(data)
-            self._message.clearNotification()
+            self._message.clear_notification()
+            try:
+                self._update_raw_data(data)
+            except (KeyError, ValueError, IndexError, UnicodeDecodeError):
+                _LOGGER.exception("Error parsing message %s", data)
+                data = b""
+
             _LOGGER.debug("%s: Received notification: %s", self.name, data)
         else:
             return
+
         new_msg = LubaMsg().parse(data)
-        if betterproto.serialized_on_wire(new_msg.net):
-            if new_msg.net.todev_ble_sync != 0 or has_field(new_msg.net.toapp_wifi_iot_status):
-                if has_field(new_msg.net.toapp_wifi_iot_status) and self._commands.get_device_product_key() == "":
+        res = betterproto2.which_one_of(new_msg, "LubaSubMsg")
+        if res[0] == "net":
+            if new_msg.net.todev_ble_sync != 0 or new_msg.net.toapp_wifi_iot_status is not None:
+                if new_msg.net.toapp_wifi_iot_status is not None and self._commands.get_device_product_key() == "":
                     self._commands.set_device_product_key(new_msg.net.toapp_wifi_iot_status.productkey)
 
-                return
+        await self._state_manager.notification(new_msg)
 
-        # may or may not be correct, some work could be done here to correctly match responses
-        if self._notify_future and not self._notify_future.done():
-            self._notify_future.set_result(data)
+        if self._execute_timed_disconnect is None:
+            await self._execute_forced_disconnect()
 
         self._reset_disconnect_timer()
-        await self._state_manager.notification(new_msg)
 
     async def _start_notify(self) -> None:
         """Start notification."""
         _LOGGER.debug("%s: Subscribe to notifications; RSSI: %s", self.name, self.rssi)
         await self._client.start_notify(self._read_char, self._notification_handler)
 
-    async def _execute_command_locked(self, key: str, command: bytes) -> bytes:
+    async def _execute_command_locked(self, key: str, command: bytes) -> None:
         """Execute command and read response."""
         assert self._client is not None
-        self._notify_future = self.loop.create_future()
-        self._key = key
         _LOGGER.debug("%s: Sending command: %s", self.name, key)
         await self._message.post_custom_data_bytes(command)
-
-        timeout = 2
-        timeout_handle = self.loop.call_at(self.loop.time() + timeout, _handle_timeout, self._notify_future)
-        timeout_expired = False
-        try:
-            notify_msg = await self._notify_future
-        except asyncio.TimeoutError:
-            timeout_expired = True
-            notify_msg = b""
-        finally:
-            if not timeout_expired:
-                timeout_handle.cancel()
-            self._notify_future = None
-
-        _LOGGER.debug("%s: Notification received: %s", self.name, notify_msg.hex())
-        return notify_msg
+        self.command_sent_time = time.time()
 
     def get_address(self) -> str:
         """Return address of device."""
-        return self._device.address
+        return self.ble_device.address
 
     def _resolve_characteristics(self, services: BleakGATTServiceCollection) -> None:
         """Resolve characteristics."""
         self._read_char = services.get_characteristic(READ_CHAR_UUID)
         if not self._read_char:
-            self._read_char = READ_CHAR_UUID
             _LOGGER.error(CharacteristicMissingError(READ_CHAR_UUID))
         self._write_char = services.get_characteristic(WRITE_CHAR_UUID)
         if not self._write_char:
-            self._write_char = WRITE_CHAR_UUID
             _LOGGER.error(CharacteristicMissingError(WRITE_CHAR_UUID))
 
     def _reset_disconnect_timer(self) -> None:
@@ -370,6 +414,7 @@ class MammotionBaseBLEDevice(MammotionBaseDevice):
             self.rssi,
         )
         self._cancel_disconnect_timer()
+        self._client = None
 
     def _disconnect_from_timer(self) -> None:
         """Disconnect from device."""
@@ -448,5 +493,5 @@ class MammotionBaseBLEDevice(MammotionBaseDevice):
             _LOGGER.debug("%s: Disconnect completed successfully", self.name)
         self._client = None
 
-    def set_disconnect_strategy(self, disconnect: bool) -> None:
+    def set_disconnect_strategy(self, *, disconnect: bool) -> None:
         self._disconnect_strategy = disconnect