401 lines
13 KiB
Python
401 lines
13 KiB
Python
#!/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
|
|
import threading
|
|
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 SwitchDpad(IntEnum):
|
|
UP = 0x00
|
|
UP_RIGHT = 0x01
|
|
RIGHT = 0x02
|
|
DOWN_RIGHT = 0x03
|
|
DOWN = 0x04
|
|
DOWN_LEFT = 0x05
|
|
LEFT = 0x06
|
|
UP_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 str_to_dpad(flags: Mapping[str, bool]) -> SwitchDpad:
|
|
"""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 SwitchDpad.UP_RIGHT
|
|
if up and left:
|
|
return SwitchDpad.UP_LEFT
|
|
if down and right:
|
|
return SwitchDpad.DOWN_RIGHT
|
|
if down and left:
|
|
return SwitchDpad.DOWN_LEFT
|
|
if up:
|
|
return SwitchDpad.UP
|
|
if down:
|
|
return SwitchDpad.DOWN
|
|
if right:
|
|
return SwitchDpad.RIGHT
|
|
if left:
|
|
return SwitchDpad.LEFT
|
|
return SwitchDpad.CENTER
|
|
|
|
|
|
@dataclass
|
|
class SwitchReport:
|
|
buttons: int = 0
|
|
hat: SwitchDpad = SwitchDpad.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_or_hat: Union[SwitchButton, SwitchDpad, int]) -> None:
|
|
"""Press one or more buttons, or set the hat if a SwitchDpad is provided."""
|
|
for item in buttons_or_hat:
|
|
if isinstance(item, SwitchDpad):
|
|
# If multiple hats are provided, the last one wins.
|
|
self.report.hat = SwitchDpad(int(item) & 0xFF)
|
|
else:
|
|
self.report.buttons |= int(item)
|
|
|
|
def release(self, *buttons_or_hat: Union[SwitchButton, SwitchDpad, int]) -> None:
|
|
"""Release one or more buttons, or center the hat if a SwitchDpad is provided."""
|
|
for item in buttons_or_hat:
|
|
if isinstance(item, SwitchDpad):
|
|
self.report.hat = SwitchDpad.CENTER
|
|
else:
|
|
self.report.buttons &= ~int(item)
|
|
|
|
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[SwitchDpad, int]) -> None:
|
|
"""Set the DPAD/hat value directly."""
|
|
self.report.hat = SwitchDpad(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 = SwitchDpad.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 = 1.0 / 500.0,
|
|
auto_send: bool = True,
|
|
) -> None:
|
|
"""
|
|
Args:
|
|
port: Serial port path (e.g., 'COM5' or '/dev/cu.usbserial-0001').
|
|
baud: UART baud rate.
|
|
send_interval: Minimum interval between sends in seconds (defaults to 500 Hz).
|
|
auto_send: If True, keep sending the current state in a background thread so the
|
|
Pico continuously sees the latest input (mirrors controller_uart_bridge).
|
|
"""
|
|
self.uart = PicoUART(port, baud)
|
|
self.state = SwitchControllerState()
|
|
self.send_interval = max(0.0, send_interval)
|
|
self._last_send = 0.0
|
|
self._auto_send = auto_send
|
|
self._stop_event = threading.Event()
|
|
self._auto_thread: Optional[threading.Thread] = None
|
|
if self._auto_send:
|
|
self._start_auto_send_thread()
|
|
|
|
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 _start_auto_send_thread(self) -> None:
|
|
"""Continuously send the current state so the Pico stays active."""
|
|
if self._auto_thread is not None:
|
|
return
|
|
sleep_time = self.send_interval if self.send_interval > 0 else 0.002
|
|
|
|
def loop() -> None:
|
|
while not self._stop_event.is_set():
|
|
self.send()
|
|
self._stop_event.wait(sleep_time)
|
|
|
|
self._auto_thread = threading.Thread(target=loop, daemon=True)
|
|
self._auto_thread.start()
|
|
|
|
def press(self, *buttons: SwitchButton | SwitchDpad | int) -> None:
|
|
"""Press buttons or set hat using SwitchButton/SwitchDpad (ints also allowed)."""
|
|
self.state.press(*buttons)
|
|
self.send()
|
|
|
|
def release(self, *buttons: SwitchButton | SwitchDpad | int) -> None:
|
|
"""Release buttons or center hat when given a SwitchDpad."""
|
|
self.state.release(*buttons)
|
|
self.send()
|
|
|
|
def set_buttons(self, buttons: Iterable[SwitchButton | int]) -> None:
|
|
self.state.set_buttons(buttons)
|
|
self.send()
|
|
|
|
def set_hat(self, hat: SwitchDpad | 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 press_for(self, duration: float, *buttons: SwitchButton | SwitchDpad | int) -> None:
|
|
"""Press buttons/hat for a duration, then release."""
|
|
self.press(*buttons)
|
|
time.sleep(max(0.0, duration))
|
|
self.release(*buttons)
|
|
|
|
def move_left_stick_for(
|
|
self, x: Union[int, float], y: Union[int, float], duration: float, neutral_after: bool = True
|
|
) -> None:
|
|
"""Move left stick for a duration, optionally returning it to neutral afterward."""
|
|
self.move_left_stick(x, y)
|
|
time.sleep(max(0.0, duration))
|
|
if neutral_after:
|
|
self.state.move_left_stick(128, 128)
|
|
self.send()
|
|
|
|
def move_right_stick_for(
|
|
self, x: Union[int, float], y: Union[int, float], duration: float, neutral_after: bool = True
|
|
) -> None:
|
|
"""Move right stick for a duration, optionally returning it to neutral afterward."""
|
|
self.move_right_stick(x, y)
|
|
time.sleep(max(0.0, duration))
|
|
if neutral_after:
|
|
self.state.move_right_stick(128, 128)
|
|
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:
|
|
if self._auto_thread:
|
|
self._stop_event.set()
|
|
self._auto_thread.join(timeout=0.5)
|
|
self._auto_thread = None
|
|
self.uart.close()
|
|
|
|
def __enter__(self) -> "SwitchUARTClient":
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
self.close()
|