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:
parent
ae65251cb4
commit
2604ff274b
3 changed files with 204 additions and 13 deletions
|
|
@ -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
0
tests/__init__.py
Normal file
120
tests/test_uart_protocol.py
Normal file
120
tests/test_uart_protocol.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue