feat(uart): add v2 protocol with IMU sample support

- Add UART_PROTOCOL_VERSION=2, IMUSample dataclass, conversion constants
- Rewrite SwitchReport.to_bytes() for versioned v2 framing with checksum
- Add ACCEL_LSB_PER_G=4096, GYRO_LSB_PER_RAD_S=818.5, MS2_PER_G constants
- Add SENSOR_ACCEL/SENSOR_GYRO SDL type constants with fallback
- Add protocol round-trip tests (8 tests, all passing)
This commit is contained in:
Joey Yakimowich-Payne 2026-03-16 11:50:25 -06:00
commit 2604ff274b
No known key found for this signature in database
GPG key ID: DDF6AF5B21B407D4
3 changed files with 204 additions and 13 deletions

View file

@ -12,6 +12,7 @@ depending on SDL. It mirrors the framing in ``switch-pico.cpp``:
from __future__ import annotations
import math
import struct
import time
import threading
@ -23,9 +24,28 @@ import serial
from serial.tools import list_ports, list_ports_common
UART_HEADER = 0xAA
UART_PROTOCOL_VERSION = 0x02
RUMBLE_HEADER = 0xBB
RUMBLE_TYPE_RUMBLE = 0x01
UART_BAUD = 921600
IMU_SAMPLES_PER_REPORT = 3
MS2_PER_G = 9.80665
RAD_TO_DEG = 180.0 / math.pi
ACCEL_LSB_PER_G = 4096.0
GYRO_LSB_PER_RAD_S = 818.5
try:
import sdl2 as _sdl2 # type: ignore[import-not-found]
_sensor_accel = getattr(_sdl2, "SDL_SENSOR_ACCEL", 1)
_sensor_gyro = getattr(_sdl2, "SDL_SENSOR_GYRO", 2)
except ImportError:
_sensor_accel = 1
_sensor_gyro = 2
SENSOR_ACCEL: int = _sensor_accel
SENSOR_GYRO: int = _sensor_gyro
class SwitchButton(IntFlag):
@ -62,9 +82,9 @@ def _is_usb_serial_path(path: str) -> bool:
"""Heuristic for USB serial path prefixes."""
lower = path.lower()
usb_prefixes = (
"/dev/ttyusb", # Linux USB serial
"/dev/ttyacm", # Linux CDC ACM
"/dev/cu.usb", # macOS cu/tty USB adapters
"/dev/ttyusb", # Linux USB serial
"/dev/ttyacm", # Linux CDC ACM
"/dev/cu.usb", # macOS cu/tty USB adapters
"/dev/tty.usb",
)
if lower.startswith(usb_prefixes):
@ -144,6 +164,7 @@ def first_serial_port(
return None
return ports[0]["device"]
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)))
@ -202,6 +223,21 @@ def str_to_dpad(flags: Mapping[str, bool]) -> SwitchDpad:
return SwitchDpad.CENTER
def compute_checksum(data: bytes) -> int:
"""Compute UART checksum as sum of bytes modulo 256."""
return sum(data) & 0xFF
@dataclass
class IMUSample:
accel_x: int = 0
accel_y: int = 0
accel_z: int = 0
gyro_x: int = 0
gyro_y: int = 0
gyro_z: int = 0
@dataclass
class SwitchReport:
buttons: int = 0
@ -210,13 +246,38 @@ class SwitchReport:
ly: int = 128
rx: int = 128
ry: int = 128
imu_samples: List[IMUSample] = field(default_factory=list)
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
"""Serialize the report into UART v2 framed packet format."""
count = min(len(self.imu_samples), IMU_SAMPLES_PER_REPORT)
payload = struct.pack(
"<HBBBBBB",
self.buttons & 0xFFFF,
int(self.hat) & 0xFF,
clamp_byte(self.lx),
clamp_byte(self.ly),
clamp_byte(self.rx),
clamp_byte(self.ry),
count,
)
for i in range(count):
sample = self.imu_samples[i]
payload += struct.pack(
"<hhhhhh",
max(-32768, min(32767, int(sample.accel_x))),
max(-32768, min(32767, int(sample.accel_y))),
max(-32768, min(32767, int(sample.accel_z))),
max(-32768, min(32767, int(sample.gyro_x))),
max(-32768, min(32767, int(sample.gyro_y))),
max(-32768, min(32767, int(sample.gyro_z))),
)
payload_len = len(payload)
frame = bytes([UART_HEADER, UART_PROTOCOL_VERSION, payload_len]) + payload
return frame + bytes([compute_checksum(frame)])
class PicoUART:
def __init__(self, port: str, baudrate: int = UART_BAUD) -> None:
@ -267,15 +328,15 @@ class PicoUART:
del self._buffer[:start]
return None
frame = self._buffer[start:start + 11]
checksum = sum(frame[:10]) & 0xFF
frame = self._buffer[start : start + 11]
checksum = compute_checksum(bytes(frame[:10]))
if frame[1] == RUMBLE_TYPE_RUMBLE and checksum == frame[10]:
payload = bytes(frame[2:10])
del self._buffer[:start + 11]
del self._buffer[: start + 11]
return payload
del self._buffer[:start + 1]
del self._buffer[: start + 1]
def close(self) -> None:
"""Close the UART connection."""
@ -434,14 +495,20 @@ class SwitchUARTClient:
self.state.move_right_stick(x, y)
self.send()
def press_for(self, duration: float, *buttons: SwitchButton | SwitchDpad | int) -> None:
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
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)
@ -451,7 +518,11 @@ class SwitchUARTClient:
self.send()
def move_right_stick_for(
self, x: Union[int, float], y: Union[int, float], duration: float, neutral_after: bool = True
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)

0
tests/__init__.py Normal file
View file

120
tests/test_uart_protocol.py Normal file
View file

@ -0,0 +1,120 @@
"""Tests for UART v2 protocol serialization in switch_pico_uart."""
import struct
import pytest
from switch_pico_bridge.switch_pico_uart import (
SwitchReport,
IMUSample,
SwitchDpad,
UART_HEADER,
UART_PROTOCOL_VERSION,
ACCEL_LSB_PER_G,
GYRO_LSB_PER_RAD_S,
MS2_PER_G,
compute_checksum,
)
def test_v2_frame_with_imu_samples():
"""V2 frame with 3 IMU samples should be 48 bytes with correct layout."""
r = SwitchReport(
buttons=0,
imu_samples=[
IMUSample(100, -200, 4096, 50, -50, 0),
IMUSample(101, -201, 4097, 51, -51, 1),
IMUSample(102, -202, 4098, 52, -52, 2),
],
)
data = r.to_bytes()
assert len(data) == 48, f"Expected 48 bytes, got {len(data)}"
assert data[0] == UART_HEADER # 0xAA
assert data[1] == UART_PROTOCOL_VERSION # 0x02
assert data[2] == 44 # payload_len
assert data[10] == 3 # imu_count
# Verify checksum
assert data[-1] == compute_checksum(data[:-1])
# Verify first sample accel_x (int16 LE at byte 11)
ax0 = struct.unpack_from("<h", data, 11)[0]
assert ax0 == 100, f"Expected accel_x=100, got {ax0}"
# Verify first sample gyro_z (int16 LE at bytes 21-22)
gz0 = struct.unpack_from("<h", data, 21)[0]
assert gz0 == 0, f"Expected gyro_z=0, got {gz0}"
def test_v2_frame_no_imu():
"""V2 frame with no IMU samples should be 12 bytes."""
r = SwitchReport(
buttons=0x0004, hat=SwitchDpad.CENTER, lx=128, ly=128, rx=128, ry=128
)
data = r.to_bytes()
assert len(data) == 12, f"Expected 12 bytes, got {len(data)}"
assert data[0] == UART_HEADER
assert data[1] == UART_PROTOCOL_VERSION
assert data[2] == 8 # payload_len
assert data[10] == 0 # imu_count
assert data[-1] == compute_checksum(data[:-1])
def test_checksum_validation():
"""Checksum should match sum of all preceding bytes & 0xFF."""
r = SwitchReport(buttons=0x0001)
data = r.to_bytes()
expected_checksum = sum(data[:-1]) & 0xFF
assert data[-1] == expected_checksum
# Corrupt a byte and verify mismatch
corrupted = bytearray(data)
corrupted[3] ^= 0xFF # flip bits in first payload byte
recalculated = sum(corrupted[:-1]) & 0xFF
assert corrupted[-1] != recalculated, "Checksum should not match corrupted data"
def test_accel_scale_gravity():
"""1G (9.80665 m/s²) should convert to ~4096 raw counts."""
# convert_accel_to_raw(9.80665) ≈ 4096
raw = int(round((MS2_PER_G / MS2_PER_G) * ACCEL_LSB_PER_G))
assert abs(raw - 4096) <= 5, f"Expected ~4096 for 1G, got {raw}"
def test_gyro_scale_one_rad():
"""1.0 rad/s should convert to ~818 raw counts."""
raw = int(round(1.0 * GYRO_LSB_PER_RAD_S))
assert abs(raw - 818) <= 5, f"Expected ~818 for 1 rad/s, got {raw}"
def test_imu_sample_dataclass():
"""IMUSample fields accept int16 range values."""
s = IMUSample(
accel_x=32767, accel_y=-32768, accel_z=0, gyro_x=100, gyro_y=-100, gyro_z=1000
)
assert s.accel_x == 32767
assert s.accel_y == -32768
assert s.gyro_z == 1000
# Values outside int16 range are clamped in to_bytes()
s2 = IMUSample(accel_x=99999)
r = SwitchReport(imu_samples=[s2])
data = r.to_bytes()
ax = struct.unpack_from("<h", data, 11)[0]
assert ax == 32767, f"Expected clamped value 32767, got {ax}"
def test_backward_compat_switch_report():
"""SwitchReport with no imu_samples produces valid v2 frame (backward compat)."""
r = SwitchReport(buttons=0x000A, lx=200, ly=50, rx=128, ry=128)
data = r.to_bytes()
assert len(data) == 12
assert data[1] == 0x02 # still v2
# Buttons at bytes 3-4
buttons = struct.unpack_from("<H", data, 3)[0]
assert buttons == 0x000A
# lx at byte 6
assert data[6] == 200
def test_max_imu_samples_capped():
"""Providing >3 IMU samples should cap at 3."""
samples = [IMUSample(i, 0, 0, 0, 0, 0) for i in range(5)]
r = SwitchReport(imu_samples=samples)
data = r.to_bytes()
assert len(data) == 48 # 3 samples, not 5
assert data[10] == 3
assert data[2] == 44 # payload_len for 3 samples