328 lines
9.8 KiB
Python
328 lines
9.8 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
|
|
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()
|