foodforthought-cli 0.2.8__py3-none-any.whl → 0.3.1__py3-none-any.whl

ate/drivers/ble_transport.py

@@ -0,0 +1,405 @@
+"""
+BLE Transport Layer for Robot Drivers.
+
+Provides a serial-compatible interface over Bluetooth Low Energy.
+Uses the bleak library for cross-platform BLE support.
+
+This enables wireless control of robots that support BLE serial
+(ESP32, HM-10, nRF52, etc.) using the same driver code as USB serial.
+
+Usage:
+    from ate.drivers.ble_transport import BLETransport
+
+    # Discover BLE devices
+    devices = await BLETransport.discover()
+
+    # Connect to a device
+    transport = BLETransport(
+        address="AA:BB:CC:DD:EE:FF",
+        service_uuid="0000ffe0-0000-1000-8000-00805f9b34fb",
+        char_uuid="0000ffe1-0000-1000-8000-00805f9b34fb",
+    )
+    await transport.connect()
+
+    # Use like serial
+    transport.write(b"command\\r\\n")
+    response = transport.read(100)
+
+    await transport.disconnect()
+"""
+
+import asyncio
+import queue
+import threading
+import time
+from dataclasses import dataclass
+from typing import Optional, List, Callable, Any
+
+try:
+    from bleak import BleakClient, BleakScanner
+    from bleak.backends.device import BLEDevice
+    HAS_BLEAK = True
+except ImportError:
+    HAS_BLEAK = False
+    BleakClient = None
+    BleakScanner = None
+    BLEDevice = None
+
+
+@dataclass
+class BLEDeviceInfo:
+    """Information about a discovered BLE device."""
+    address: str
+    name: str
+    rssi: int
+    services: List[str]
+
+    def __str__(self) -> str:
+        return f"{self.name or 'Unknown'} ({self.address}) RSSI: {self.rssi}"
+
+
+class BLETransport:
+    """
+    BLE transport that provides a serial-compatible interface.
+
+    This allows robot drivers to use BLE with minimal changes.
+    The interface mimics pyserial's Serial class.
+    """
+
+    # Default ESP32/HM-10 BLE serial UUIDs
+    DEFAULT_SERVICE_UUID = "0000ffe0-0000-1000-8000-00805f9b34fb"
+    DEFAULT_WRITE_CHAR_UUID = "0000ffe1-0000-1000-8000-00805f9b34fb"  # Write characteristic
+    DEFAULT_NOTIFY_CHAR_UUID = "0000ffe2-0000-1000-8000-00805f9b34fb"  # Notify characteristic
+    # Legacy single-char UUID for devices that use same char for both
+    DEFAULT_CHAR_UUID = "0000ffe1-0000-1000-8000-00805f9b34fb"
+
+    def __init__(
+        self,
+        address: str,
+        service_uuid: str = None,
+        char_uuid: str = None,
+        write_char_uuid: str = None,
+        notify_char_uuid: str = None,
+        timeout: float = 2.0,
+    ):
+        """
+        Initialize BLE transport.
+
+        Args:
+            address: BLE device address (MAC on Windows/Linux, UUID on macOS)
+            service_uuid: BLE service UUID (default: HM-10 compatible)
+            char_uuid: BLE characteristic UUID for read/write (legacy, single char)
+            write_char_uuid: BLE characteristic UUID for writing (ffe1)
+            notify_char_uuid: BLE characteristic UUID for notifications (ffe2)
+            timeout: Read timeout in seconds
+        """
+        if not HAS_BLEAK:
+            raise ImportError("bleak is required for BLE. Install with: pip install bleak")
+
+        self.address = address
+        self.service_uuid = service_uuid or self.DEFAULT_SERVICE_UUID
+
+        # Support both legacy single-char and separate write/notify chars
+        if write_char_uuid or notify_char_uuid:
+            self.write_char_uuid = write_char_uuid or self.DEFAULT_WRITE_CHAR_UUID
+            self.notify_char_uuid = notify_char_uuid or self.DEFAULT_NOTIFY_CHAR_UUID
+        elif char_uuid:
+            # Legacy: same char for both
+            self.write_char_uuid = char_uuid
+            self.notify_char_uuid = char_uuid
+        else:
+            # Default: use separate chars (MechDog style)
+            self.write_char_uuid = self.DEFAULT_WRITE_CHAR_UUID
+            self.notify_char_uuid = self.DEFAULT_NOTIFY_CHAR_UUID
+
+        # Keep char_uuid for backwards compatibility
+        self.char_uuid = self.write_char_uuid
+        self.timeout = timeout
+
+        self._client: Optional[BleakClient] = None
+        self._connected = False
+        self._rx_buffer = queue.Queue()
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._thread: Optional[threading.Thread] = None
+
+    @staticmethod
+    async def discover(
+        timeout: float = 10.0,
+        name_filter: str = None,
+    ) -> List[BLEDeviceInfo]:
+        """
+        Discover BLE devices.
+
+        Args:
+            timeout: Scan duration in seconds
+            name_filter: Only return devices whose name contains this string
+
+        Returns:
+            List of discovered BLE devices
+        """
+        if not HAS_BLEAK:
+            raise ImportError("bleak is required for BLE. Install with: pip install bleak")
+
+        devices = await BleakScanner.discover(timeout=timeout)
+
+        result = []
+        for d in devices:
+            if name_filter and (not d.name or name_filter.lower() not in d.name.lower()):
+                continue
+
+            info = BLEDeviceInfo(
+                address=d.address,
+                name=d.name or "Unknown",
+                rssi=d.rssi if hasattr(d, 'rssi') else -100,
+                services=[],  # Would need to connect to enumerate
+            )
+            result.append(info)
+
+        # Sort by signal strength
+        result.sort(key=lambda x: x.rssi, reverse=True)
+        return result
+
+    @staticmethod
+    def discover_sync(timeout: float = 10.0, name_filter: str = None) -> List[BLEDeviceInfo]:
+        """Synchronous wrapper for discover()."""
+        return asyncio.run(BLETransport.discover(timeout, name_filter))
+
+    def _notification_handler(self, sender: int, data: bytearray):
+        """Handle incoming BLE notifications."""
+        self._rx_buffer.put(bytes(data))
+
+    async def _connect_async(self) -> bool:
+        """Async connect implementation."""
+        try:
+            self._client = BleakClient(self.address)
+            await self._client.connect()
+
+            # Subscribe to notifications on the NOTIFY characteristic (ffe2)
+            await self._client.start_notify(self.notify_char_uuid, self._notification_handler)
+
+            self._connected = True
+            return True
+        except Exception as e:
+            print(f"BLE connect error: {e}")
+            self._connected = False
+            return False
+
+    async def _disconnect_async(self):
+        """Async disconnect implementation."""
+        if self._client and self._connected:
+            try:
+                await self._client.stop_notify(self.notify_char_uuid)
+                await self._client.disconnect()
+            except Exception:
+                pass
+        self._connected = False
+        self._client = None
+
+    async def _write_async(self, data: bytes):
+        """Async write implementation."""
+        if not self._connected or not self._client:
+            raise IOError("Not connected")
+
+        # Split into chunks if needed (BLE has MTU limits)
+        chunk_size = 20  # Safe default, could be negotiated higher
+        for i in range(0, len(data), chunk_size):
+            chunk = data[i:i + chunk_size]
+            # Write to the WRITE characteristic (ffe1), without response for speed
+            await self._client.write_gatt_char(self.write_char_uuid, chunk, response=False)
+
+    def _run_async_loop(self):
+        """Run the async event loop in a background thread."""
+        self._loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def _run_coroutine(self, coro) -> Any:
+        """Run a coroutine from sync code."""
+        if self._loop is None:
+            # Start the async loop in a background thread
+            self._thread = threading.Thread(target=self._run_async_loop, daemon=True)
+            self._thread.start()
+            time.sleep(0.1)  # Give the loop time to start
+
+        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        return future.result(timeout=self.timeout * 2)
+
+    def connect(self) -> bool:
+        """
+        Connect to the BLE device.
+
+        Returns:
+            True if connected successfully
+        """
+        return self._run_coroutine(self._connect_async())
+
+    def disconnect(self):
+        """Disconnect from the BLE device."""
+        if self._loop:
+            self._run_coroutine(self._disconnect_async())
+
+        if self._loop:
+            self._loop.call_soon_threadsafe(self._loop.stop)
+        if self._thread:
+            self._thread.join(timeout=1.0)
+
+    def write(self, data: bytes) -> int:
+        """
+        Write data to the BLE device.
+
+        Args:
+            data: Bytes to send
+
+        Returns:
+            Number of bytes written
+        """
+        self._run_coroutine(self._write_async(data))
+        return len(data)
+
+    def read(self, size: int = 1) -> bytes:
+        """
+        Read data from the BLE device.
+
+        Args:
+            size: Maximum bytes to read
+
+        Returns:
+            Received bytes (may be less than size)
+        """
+        result = b""
+        deadline = time.time() + self.timeout
+
+        while len(result) < size and time.time() < deadline:
+            try:
+                chunk = self._rx_buffer.get(timeout=0.1)
+                result += chunk
+            except queue.Empty:
+                if result:  # Got some data, return it
+                    break
+
+        return result[:size]
+
+    def read_until(self, terminator: bytes = b"\n") -> bytes:
+        """
+        Read until terminator is found.
+
+        Args:
+            terminator: Bytes to read until
+
+        Returns:
+            Data including terminator
+        """
+        result = b""
+        deadline = time.time() + self.timeout
+
+        while time.time() < deadline:
+            try:
+                chunk = self._rx_buffer.get(timeout=0.1)
+                result += chunk
+                if terminator in result:
+                    break
+            except queue.Empty:
+                continue
+
+        return result
+
+    def readline(self) -> bytes:
+        """Read a line (until newline)."""
+        return self.read_until(b"\n")
+
+    def flush(self):
+        """Flush buffers (no-op for BLE)."""
+        pass
+
+    def reset_input_buffer(self):
+        """Clear the input buffer."""
+        while not self._rx_buffer.empty():
+            try:
+                self._rx_buffer.get_nowait()
+            except queue.Empty:
+                break
+
+    def reset_output_buffer(self):
+        """Clear the output buffer (no-op for BLE)."""
+        pass
+
+    @property
+    def is_open(self) -> bool:
+        """Check if connection is open."""
+        return self._connected
+
+    @property
+    def in_waiting(self) -> int:
+        """Number of bytes in input buffer."""
+        return self._rx_buffer.qsize() * 20  # Approximate
+
+    def close(self):
+        """Close the connection."""
+        self.disconnect()
+
+    def __enter__(self):
+        self.connect()
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+
+def discover_ble_robots(
+    timeout: float = 10.0,
+    name_patterns: List[str] = None,
+) -> List[BLEDeviceInfo]:
+    """
+    Discover BLE robots by scanning for known device names.
+
+    Args:
+        timeout: Scan duration
+        name_patterns: Device name patterns to match (e.g., ["MechDog", "ESP32"])
+
+    Returns:
+        List of discovered robot devices
+    """
+    if not HAS_BLEAK:
+        print("Warning: bleak not installed. Run: pip install bleak")
+        return []
+
+    if name_patterns is None:
+        name_patterns = [
+            "MechDog",
+            "HiWonder",
+            "ESP32",
+            "ESP-",
+            "BT",
+            "HC-",
+            "HM-",
+        ]
+
+    all_devices = BLETransport.discover_sync(timeout=timeout)
+
+    # Filter by name patterns
+    robots = []
+    for device in all_devices:
+        if not device.name:
+            continue
+        for pattern in name_patterns:
+            if pattern.lower() in device.name.lower():
+                robots.append(device)
+                break
+
+    return robots
+
+
+if __name__ == "__main__":
+    # Test discovery
+    print("Scanning for BLE devices...")
+    devices = BLETransport.discover_sync(timeout=5.0)
+
+    print(f"\nFound {len(devices)} devices:")
+    for d in devices:
+        print(f"  {d}")
+
+    # Filter for potential robots
+    print("\nPotential robots:")
+    robots = discover_ble_robots(timeout=5.0)
+    for r in robots:
+        print(f"  {r}")