Refactor into lib

switch_pico_uart.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""
+Lightweight helpers for talking to the switch-pico firmware over UART.
+
+This module exposes the raw report structure plus a small convenience wrapper
+so other scripts can do things like "press a button" or "move a stick" without
+depending on SDL. It mirrors the framing in ``switch-pico.cpp``:
+
+  Host -> Pico : 0xAA, buttons (LE16), hat, lx, ly, rx, ry
+  Pico -> Host : 0xBB, 0x01, 8 rumble bytes, checksum (sum of first 10 bytes)
+"""
+
+from __future__ import annotations
+
+import struct
+import time
+from dataclasses import dataclass, field
+from enum import IntEnum, IntFlag
+from typing import Iterable, Mapping, Optional, Tuple, Union
+
+import serial
+
+UART_HEADER = 0xAA
+RUMBLE_HEADER = 0xBB
+RUMBLE_TYPE_RUMBLE = 0x01
+UART_BAUD = 921600
+
+
+class SwitchButton(IntFlag):
+    # Mirrors the masks defined in switch_pro_descriptors.h
+    Y = 1 << 0
+    B = 1 << 1
+    A = 1 << 2
+    X = 1 << 3
+    L = 1 << 4
+    R = 1 << 5
+    ZL = 1 << 6
+    ZR = 1 << 7
+    MINUS = 1 << 8
+    PLUS = 1 << 9
+    LCLICK = 1 << 10
+    RCLICK = 1 << 11
+    HOME = 1 << 12
+    CAPTURE = 1 << 13
+
+
+class SwitchHat(IntEnum):
+    TOP = 0x00
+    TOP_RIGHT = 0x01
+    RIGHT = 0x02
+    BOTTOM_RIGHT = 0x03
+    BOTTOM = 0x04
+    BOTTOM_LEFT = 0x05
+    LEFT = 0x06
+    TOP_LEFT = 0x07
+    CENTER = 0x08
+
+
+def clamp_byte(value: Union[int, float]) -> int:
+    """Clamp a numeric value to the 0-255 byte range."""
+    return max(0, min(255, int(value)))
+
+
+def normalize_stick_value(value: Union[int, float]) -> int:
+    """
+    Convert a normalized float (-1..1) or raw byte (0..255) to the stick range.
+
+    Floats are treated as -1.0 = full negative deflection, 0.0 = center,
+    1.0 = full positive deflection. Integers are assumed to already be in the
+    0-255 range.
+    """
+    if isinstance(value, float):
+        value = max(-1.0, min(1.0, value))
+        value = int(round((value + 1.0) * 255 / 2.0))
+    return clamp_byte(value)
+
+
+def axis_to_stick(value: int, deadzone: int) -> int:
+    """Convert a signed axis value to 0-255 stick range with deadzone."""
+    if abs(value) < deadzone:
+        value = 0
+    scaled = int((value + 32768) * 255 / 65535)
+    return clamp_byte(scaled)
+
+
+def trigger_to_button(value: int, threshold: int) -> bool:
+    """Return True if analog trigger crosses digital threshold."""
+    return value >= threshold
+
+
+def dpad_to_hat(flags: Mapping[str, bool]) -> SwitchHat:
+    """Translate DPAD button flags into a Switch hat value."""
+    up = flags.get("up", False)
+    down = flags.get("down", False)
+    left = flags.get("left", False)
+    right = flags.get("right", False)
+
+    if up and right:
+        return SwitchHat.TOP_RIGHT
+    if up and left:
+        return SwitchHat.TOP_LEFT
+    if down and right:
+        return SwitchHat.BOTTOM_RIGHT
+    if down and left:
+        return SwitchHat.BOTTOM_LEFT
+    if up:
+        return SwitchHat.TOP
+    if down:
+        return SwitchHat.BOTTOM
+    if right:
+        return SwitchHat.RIGHT
+    if left:
+        return SwitchHat.LEFT
+    return SwitchHat.CENTER
+
+
+@dataclass
+class SwitchReport:
+    buttons: int = 0
+    hat: SwitchHat = SwitchHat.CENTER
+    lx: int = 128
+    ly: int = 128
+    rx: int = 128
+    ry: int = 128
+
+    def to_bytes(self) -> bytes:
+        """Serialize the report into the UART packet format."""
+        return struct.pack(
+            "<BHBBBBB", UART_HEADER, self.buttons & 0xFFFF, self.hat & 0xFF, self.lx, self.ly, self.rx, self.ry
+        )
+
+
+class PicoUART:
+    def __init__(self, port: str, baudrate: int = UART_BAUD) -> None:
+        """Open a UART connection to the Pico with non-blocking IO."""
+        self.serial = serial.Serial(
+            port=port,
+            baudrate=baudrate,
+            bytesize=serial.EIGHTBITS,
+            stopbits=serial.STOPBITS_ONE,
+            parity=serial.PARITY_NONE,
+            timeout=0.0,
+            write_timeout=0.0,
+            xonxoff=False,
+            rtscts=False,
+            dsrdtr=False,
+        )
+        self._buffer = bytearray()
+
+    def send_report(self, report: SwitchReport) -> None:
+        """Send a controller report to the Pico."""
+        self.serial.write(report.to_bytes())
+
+    def read_rumble_payload(self) -> Optional[bytes]:
+        """
+        Drain available UART bytes into an internal buffer, then extract one rumble frame.
+
+        Frame format:
+          0: 0xBB (RUMBLE_HEADER)
+          1: type (0x01 for rumble)
+          2-9: 8-byte rumble payload
+          10: checksum (sum of first 10 bytes) & 0xFF
+        """
+        waiting = self.serial.in_waiting
+        if waiting:
+            self._buffer.extend(self.serial.read(waiting))
+
+        while True:
+            if not self._buffer:
+                return None
+
+            start = self._buffer.find(bytes([RUMBLE_HEADER]))
+            if start < 0:
+                self._buffer.clear()
+                return None
+
+            if len(self._buffer) - start < 11:
+                if start > 0:
+                    del self._buffer[:start]
+                return None
+
+            frame = self._buffer[start:start + 11]
+            checksum = sum(frame[:10]) & 0xFF
+
+            if frame[1] == RUMBLE_TYPE_RUMBLE and checksum == frame[10]:
+                payload = bytes(frame[2:10])
+                del self._buffer[:start + 11]
+                return payload
+
+            del self._buffer[:start + 1]
+
+    def close(self) -> None:
+        """Close the UART connection."""
+        self.serial.close()
+
+
+def decode_rumble(payload: bytes) -> Tuple[float, float]:
+    """Return normalized rumble amplitudes (0.0-1.0) for left/right."""
+    if len(payload) < 8:
+        return 0.0, 0.0
+    if payload == b"\x00\x01\x40\x40\x00\x01\x40\x40":
+        return 0.0, 0.0
+    right_raw = ((payload[1] & 0x03) << 8) | payload[0]
+    left_raw = ((payload[5] & 0x03) << 8) | payload[4]
+    if left_raw < 8 and right_raw < 8:
+        return 0.0, 0.0
+    left = min(max(left_raw / 1023.0, 0.0), 1.0)
+    right = min(max(right_raw / 1023.0, 0.0), 1.0)
+    return left, right
+
+
+@dataclass
+class SwitchControllerState:
+    """Mutable controller state with helpers for building reports."""
+
+    report: SwitchReport = field(default_factory=SwitchReport)
+
+    def press(self, *buttons: Union[SwitchButton, int]) -> None:
+        """Set one or more buttons as pressed."""
+        for button in buttons:
+            self.report.buttons |= int(button)
+
+    def release(self, *buttons: Union[SwitchButton, int]) -> None:
+        """Release one or more buttons."""
+        for button in buttons:
+            self.report.buttons &= ~int(button)
+
+    def set_buttons(self, buttons: Iterable[Union[SwitchButton, int]]) -> None:
+        """Replace the current button bitmask with the provided buttons."""
+        self.report.buttons = 0
+        self.press(*buttons)
+
+    def set_hat(self, hat: Union[SwitchHat, int]) -> None:
+        """Set the DPAD/hat value directly."""
+        self.report.hat = int(hat) & 0xFF
+
+    def move_left_stick(self, x: Union[int, float], y: Union[int, float]) -> None:
+        """Move the left stick using normalized floats (-1..1) or raw bytes (0-255)."""
+        self.report.lx = normalize_stick_value(x)
+        self.report.ly = normalize_stick_value(y)
+
+    def move_right_stick(self, x: Union[int, float], y: Union[int, float]) -> None:
+        """Move the right stick using normalized floats (-1..1) or raw bytes (0-255)."""
+        self.report.rx = normalize_stick_value(x)
+        self.report.ry = normalize_stick_value(y)
+
+    def neutral(self) -> None:
+        """Clear all input back to the neutral controller state."""
+        self.report.buttons = 0
+        self.report.hat = SwitchHat.CENTER
+        self.report.lx = 128
+        self.report.ly = 128
+        self.report.rx = 128
+        self.report.ry = 128
+
+
+class SwitchUARTClient:
+    """
+    High-level helper to send controller actions to the Pico and poll rumble.
+
+    Example:
+        with SwitchUARTClient("/dev/cu.usbserial-0001") as client:
+            client.press(SwitchButton.A)
+            time.sleep(0.1)
+            client.release(SwitchButton.A)
+            client.move_left_stick(0.0, -1.0)  # push up
+    """
+
+    def __init__(self, port: str, baud: int = UART_BAUD, send_interval: float = 0.0) -> None:
+        self.uart = PicoUART(port, baud)
+        self.state = SwitchControllerState()
+        self.send_interval = max(0.0, send_interval)
+        self._last_send = 0.0
+
+    def send(self) -> None:
+        """Send the current state to the Pico, throttled by send_interval if set."""
+        now = time.monotonic()
+        if self.send_interval and (now - self._last_send) < self.send_interval:
+            return
+        self.uart.send_report(self.state.report)
+        self._last_send = now
+
+    def press(self, *buttons: int) -> None:
+        self.state.press(*buttons)
+        self.send()
+
+    def release(self, *buttons: int) -> None:
+        self.state.release(*buttons)
+        self.send()
+
+    def set_buttons(self, buttons: Iterable[int]) -> None:
+        self.state.set_buttons(buttons)
+        self.send()
+
+    def set_hat(self, hat: int) -> None:
+        self.state.set_hat(hat)
+        self.send()
+
+    def move_left_stick(self, x: Union[int, float], y: Union[int, float]) -> None:
+        self.state.move_left_stick(x, y)
+        self.send()
+
+    def move_right_stick(self, x: Union[int, float], y: Union[int, float]) -> None:
+        self.state.move_right_stick(x, y)
+        self.send()
+
+    def neutral(self) -> None:
+        self.state.neutral()
+        self.send()
+
+    def poll_rumble(self) -> Optional[Tuple[float, float]]:
+        """
+        Poll for the latest rumble payload and return normalized amplitudes.
+        Returns None if no rumble frame was available.
+        """
+        payload = self.uart.read_rumble_payload()
+        if payload:
+            return decode_rumble(payload)
+        return None
+
+    def close(self) -> None:
+        self.uart.close()
+
+    def __enter__(self) -> "SwitchUARTClient":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()