foodforthought-cli 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/transports/__init__.py

@@ -0,0 +1,121 @@
+"""
+Transport Abstraction Layer for Robot Communication.
+
+This module provides a unified interface for communicating with robots
+across different physical transports (BLE, USB Serial, WiFi, etc.).
+
+The key insight: Users shouldn't need to know HOW a robot communicates,
+only WHAT they want it to do.
+
+Example:
+    from ate.transports import connect
+
+    # Auto-discovers available transports
+    robot = await connect("mechdog")
+
+    # These work regardless of underlying transport
+    await robot.walk_forward()
+    await robot.gripper_open()  # Automatically routes to USB if BLE can't do it
+"""
+
+from typing import Optional
+
+from .base import (
+    RobotTransport,
+    TransportCapability,
+    TransportState,
+    ConnectionInfo,
+    CommandResult,
+    HiWonderProtocol,
+)
+from .ble import BLETransport
+from .serial import SerialTransport
+from .hybrid import HybridTransport
+
+
+async def connect(
+    address: Optional[str] = None,
+    name: Optional[str] = None,
+    transport_type: str = "auto"
+) -> HybridTransport:
+    """
+    Connect to a robot with automatic transport detection.
+
+    This is the main entry point for the transport layer.
+    It discovers available transports and creates a hybrid connection.
+
+    Args:
+        address: Specific address to connect to (BLE MAC/UUID or serial port)
+        name: Robot name hint for discovery (e.g., "mechdog")
+        transport_type: "auto", "ble", "serial", or "hybrid"
+
+    Returns:
+        Connected HybridTransport instance
+
+    Example:
+        # Auto-discover and connect
+        robot = await connect(name="mechdog")
+        await robot.walk_forward()
+
+        # Connect to specific BLE device
+        robot = await connect(address="FAF198B6-D9A4-CC27-7BBA-3159F63A61A9")
+
+        # Connect to specific serial port
+        robot = await connect(address="/dev/cu.usbmodem1234")
+    """
+    robot = HybridTransport()
+
+    if address:
+        # Direct address provided
+        await robot.connect(address, name=name)
+    else:
+        # Discover and connect to first matching device
+        connections = await HybridTransport.discover(timeout=5.0)
+
+        ble_device = None
+        serial_device = None
+
+        for conn in connections:
+            if name and conn.name and name.lower() not in conn.name.lower():
+                continue
+
+            if conn.transport_type == "ble" and not ble_device:
+                ble_device = conn
+            elif conn.transport_type == "serial" and not serial_device:
+                serial_device = conn
+
+        if ble_device:
+            await robot.connect_ble(ble_device.address)
+        if serial_device:
+            await robot.connect_serial(serial_device.address)
+
+    return robot
+
+
+async def discover(timeout: float = 5.0) -> list[ConnectionInfo]:
+    """
+    Discover all available robots across all transports.
+
+    Returns:
+        List of discovered robot connections
+    """
+    return await HybridTransport.discover(timeout)
+
+
+__all__ = [
+    # Core types
+    "RobotTransport",
+    "TransportCapability",
+    "TransportState",
+    "ConnectionInfo",
+    "CommandResult",
+    # Protocol adapters
+    "HiWonderProtocol",
+    # Transport implementations
+    "BLETransport",
+    "SerialTransport",
+    "HybridTransport",
+    # Convenience functions
+    "connect",
+    "discover",
+]

ate/transports/base.py

@@ -0,0 +1,394 @@
+"""
+Base Transport Interface.
+
+Defines the abstract interface that all robot transports must implement.
+This allows the rest of the codebase to work with robots without knowing
+the underlying communication protocol.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum, auto
+from typing import Optional, Callable, Awaitable, Any
+import asyncio
+
+
+class TransportState(Enum):
+    """Current state of a transport connection."""
+    DISCONNECTED = auto()
+    CONNECTING = auto()
+    CONNECTED = auto()
+    ERROR = auto()
+
+
+class TransportCapability(Enum):
+    """
+    Capabilities that a transport may or may not support.
+
+    This is CRITICAL for the abstraction layer - we discovered that:
+    - BLE can do locomotion but NOT arm control on MechDog
+    - USB Serial can do arm control but connection is finicky
+    - Different firmware versions support different commands
+
+    By probing capabilities, we can automatically route commands
+    to the right transport.
+    """
+    # Locomotion
+    WALK = auto()
+    TURN = auto()
+    STAND = auto()
+    SIT = auto()
+    POSTURE = auto()
+
+    # Manipulation
+    ARM_POSITION = auto()
+    GRIPPER = auto()
+
+    # Sensors
+    ULTRASONIC = auto()
+    IMU = auto()
+    BATTERY = auto()
+
+    # Actions (pre-programmed sequences)
+    ACTION_GROUPS = auto()
+
+    # LED/Visual
+    RGB_LED = auto()
+
+    # Low-level
+    RAW_SERVO = auto()
+
+    # Meta
+    BIDIRECTIONAL = auto()  # Can receive responses
+    STREAMING = auto()       # Can stream continuous data
+
+
+@dataclass
+class ConnectionInfo:
+    """Information about a discovered or connected transport."""
+    transport_type: str  # "ble", "serial", "wifi"
+    address: str         # MAC address, port, or IP
+    name: Optional[str] = None
+    rssi: Optional[int] = None  # Signal strength for wireless
+    metadata: dict = field(default_factory=dict)
+
+
+@dataclass
+class CommandResult:
+    """Result of sending a command."""
+    success: bool
+    response: Optional[str] = None
+    raw_data: Optional[bytes] = None
+    error: Optional[str] = None
+    latency_ms: Optional[float] = None
+
+
+class RobotTransport(ABC):
+    """
+    Abstract base class for robot communication transports.
+
+    Every transport (BLE, Serial, WiFi) implements this interface,
+    allowing higher-level code to work with any robot without
+    caring about the underlying protocol.
+
+    Key design decisions:
+    1. Async-first: Robot communication is inherently async
+    2. Capability-based: Not all transports support all features
+    3. Response handling: Can register callbacks for async responses
+    4. Connection lifecycle: Connect/disconnect/reconnect handling
+    """
+
+    def __init__(self):
+        self._state = TransportState.DISCONNECTED
+        self._capabilities: set[TransportCapability] = set()
+        self._response_handlers: list[Callable[[bytes], Awaitable[None]]] = []
+        self._connection_info: Optional[ConnectionInfo] = None
+
+    @property
+    def state(self) -> TransportState:
+        """Current connection state."""
+        return self._state
+
+    @property
+    def is_connected(self) -> bool:
+        """Whether transport is currently connected."""
+        return self._state == TransportState.CONNECTED
+
+    @property
+    def capabilities(self) -> set[TransportCapability]:
+        """Set of capabilities this transport supports."""
+        return self._capabilities.copy()
+
+    @property
+    def connection_info(self) -> Optional[ConnectionInfo]:
+        """Information about the current connection."""
+        return self._connection_info
+
+    def has_capability(self, cap: TransportCapability) -> bool:
+        """Check if this transport supports a specific capability."""
+        return cap in self._capabilities
+
+    # ========== Abstract Methods ==========
+
+    @abstractmethod
+    async def connect(self, address: str, **kwargs) -> bool:
+        """
+        Connect to a robot at the given address.
+
+        Args:
+            address: Transport-specific address (MAC, port, IP)
+            **kwargs: Transport-specific options
+
+        Returns:
+            True if connection successful
+        """
+        pass
+
+    @abstractmethod
+    async def disconnect(self) -> None:
+        """Disconnect from the robot."""
+        pass
+
+    @abstractmethod
+    async def send(self, data: bytes) -> CommandResult:
+        """
+        Send raw data to the robot.
+
+        Args:
+            data: Raw bytes to send
+
+        Returns:
+            CommandResult with success status and any response
+        """
+        pass
+
+    @abstractmethod
+    async def send_command(self, command: str) -> CommandResult:
+        """
+        Send a formatted command string to the robot.
+
+        This handles protocol-specific formatting (e.g., HiWonder CMD|...|$)
+
+        Args:
+            command: Command string in transport's native format
+
+        Returns:
+            CommandResult with success status and any response
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    async def discover(cls, timeout: float = 5.0) -> list[ConnectionInfo]:
+        """
+        Discover available robots using this transport.
+
+        Args:
+            timeout: How long to scan in seconds
+
+        Returns:
+            List of discovered connection info
+        """
+        pass
+
+    @abstractmethod
+    async def probe_capabilities(self) -> set[TransportCapability]:
+        """
+        Probe the connected robot to discover what it can do.
+
+        This is where we send test commands and see what works.
+        Essential for handling firmware fragmentation.
+
+        Returns:
+            Set of capabilities that work on this connection
+        """
+        pass
+
+    # ========== Response Handling ==========
+
+    def add_response_handler(
+        self,
+        handler: Callable[[bytes], Awaitable[None]]
+    ) -> None:
+        """Register a callback for async responses from robot."""
+        self._response_handlers.append(handler)
+
+    def remove_response_handler(
+        self,
+        handler: Callable[[bytes], Awaitable[None]]
+    ) -> None:
+        """Remove a response handler."""
+        if handler in self._response_handlers:
+            self._response_handlers.remove(handler)
+
+    async def _notify_handlers(self, data: bytes) -> None:
+        """Notify all registered handlers of received data."""
+        for handler in self._response_handlers:
+            try:
+                await handler(data)
+            except Exception as e:
+                # Don't let one handler crash others
+                print(f"Handler error: {e}")
+
+    # ========== Context Manager ==========
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit - ensures cleanup."""
+        await self.disconnect()
+        return False
+
+
+class ProtocolAdapter(ABC):
+    """
+    Adapts high-level commands to protocol-specific format.
+
+    Different robots use different command formats:
+    - HiWonder: CMD|function|sub|data|$
+    - Unitree: Binary protocol
+    - Some: JSON over WebSocket
+
+    The adapter translates between canonical commands and
+    protocol-specific formats.
+    """
+
+    @abstractmethod
+    def encode_command(self, command: str, *args) -> bytes:
+        """
+        Encode a canonical command into protocol-specific bytes.
+
+        Args:
+            command: Canonical command name (e.g., "walk_forward")
+            *args: Command arguments
+
+        Returns:
+            Protocol-specific bytes to send
+        """
+        pass
+
+    @abstractmethod
+    def decode_response(self, data: bytes) -> dict[str, Any]:
+        """
+        Decode protocol-specific response into canonical format.
+
+        Args:
+            data: Raw bytes received from robot
+
+        Returns:
+            Decoded response as a dictionary
+        """
+        pass
+
+
+class HiWonderProtocol(ProtocolAdapter):
+    """
+    Protocol adapter for HiWonder robots (MechDog, etc).
+
+    Command format: CMD|function|subfunction|data...|$
+    Response format: CMD|function|data...|$
+    """
+
+    # Canonical command -> (function, subfunction, args_formatter)
+    COMMAND_MAP = {
+        # Posture commands (Function 1)
+        "posture_normal": (1, 0, None),
+        "posture_high": (1, 1, None),
+        "posture_low": (1, 2, None),
+        "turn_left": (1, 3, None),
+        "turn_right": (1, 4, None),
+        "posture_reset": (1, 5, None),
+        "sit": (1, 7, None),
+        "lie_down": (1, 8, None),
+        "low_crawl": (1, 9, None),
+        "high_crawl": (1, 10, None),
+        "twist": (1, 11, None),
+        "wave": (1, 12, None),
+        "handshake": (1, 13, None),
+
+        # Action groups (Function 2)
+        "action_group": (2, 1, lambda group_id: [group_id]),
+
+        # Movement (Function 3)
+        "move": (3, None, lambda x, y, yaw: [x, y, yaw]),
+        "walk_forward": (3, None, lambda: [100, 0, 0]),
+        "walk_backward": (3, None, lambda: [-100, 0, 0]),
+        "strafe_left": (3, None, lambda: [0, -100, 0]),
+        "strafe_right": (3, None, lambda: [0, 100, 0]),
+        "rotate_left": (3, None, lambda: [0, 0, 50]),
+        "rotate_right": (3, None, lambda: [0, 0, -50]),
+        "stop": (3, None, lambda: [0, 0, 0]),
+
+        # Sensors (Function 4)
+        "ultrasonic_query": (4, 1, None),
+        "rgb_led": (4, 3, lambda r, g, b: [r, g, b]),
+
+        # Battery (Function 6)
+        "battery_query": (6, None, None),
+
+        # Arm control (Function 7) - May not work on all firmware
+        "arm_action_1": (7, 1, None),
+        "arm_action_2": (7, 2, None),
+        "arm_joint_a_plus": (7, 3, lambda: [1]),
+        "arm_joint_a_minus": (7, 3, lambda: [-1]),
+        "arm_joint_b_plus": (7, 4, lambda: [1]),
+        "arm_joint_b_minus": (7, 4, lambda: [-1]),
+        "gripper_open": (7, 6, None),
+        "gripper_close": (7, 7, None),
+    }
+
+    def encode_command(self, command: str, *args) -> bytes:
+        """Encode canonical command to HiWonder format."""
+        if command not in self.COMMAND_MAP:
+            raise ValueError(f"Unknown command: {command}")
+
+        func, sub, args_formatter = self.COMMAND_MAP[command]
+
+        parts = ["CMD", str(func)]
+
+        if sub is not None:
+            parts.append(str(sub))
+
+        if args_formatter:
+            formatted_args = args_formatter(*args) if args else args_formatter()
+            parts.extend(str(a) for a in formatted_args)
+
+        parts.append("$")
+        return "|".join(parts).encode('ascii')
+
+    def decode_response(self, data: bytes) -> dict[str, Any]:
+        """Decode HiWonder response."""
+        try:
+            text = data.decode('ascii').strip()
+            if not text.startswith("CMD|"):
+                return {"raw": text}
+
+            parts = text.split("|")
+            func = int(parts[1])
+
+            result = {
+                "function": func,
+                "raw": text,
+                "parts": parts,
+            }
+
+            # Parse known response types
+            if func == 4 and len(parts) >= 3:
+                # Ultrasonic response: CMD|4|<mm>|$
+                try:
+                    result["ultrasonic_mm"] = int(parts[2])
+                except ValueError:
+                    pass
+            elif func == 6 and len(parts) >= 3:
+                # Battery response: CMD|6|<voltage>|$ (voltage in 0.01V)
+                try:
+                    result["battery_voltage"] = float(parts[2])
+                except ValueError:
+                    pass
+
+            return result
+
+        except Exception as e:
+            return {"error": str(e), "raw_bytes": data}