vee
131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
"""Optional MQTT 3.1.1 client used by the browser workstation."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import threading
|
||
import uuid
|
||
from collections.abc import Callable
|
||
from typing import Any
|
||
|
||
StateCallback = Callable[[str, str], None]
|
||
PayloadCallback = Callable[[str, str], None]
|
||
|
||
|
||
class MqttSeatPoseClient:
|
||
def __init__(
|
||
self,
|
||
host: str,
|
||
port: int,
|
||
topic: str,
|
||
on_state: StateCallback,
|
||
on_payload: PayloadCallback,
|
||
) -> None:
|
||
self.host = host
|
||
self.port = port
|
||
self.topic = topic
|
||
self.test_topic = f"n7calibration/seatpose_test/{uuid.uuid4().hex[:12]}"
|
||
self.on_state = on_state
|
||
self.on_payload = on_payload
|
||
self._client: Any = None
|
||
self._lock = threading.Lock()
|
||
|
||
def connect(self) -> None:
|
||
try:
|
||
import paho.mqtt.client as mqtt # type: ignore
|
||
except ImportError as error:
|
||
raise RuntimeError("缺少 paho-mqtt；请安装项目的 ir 可选依赖") from error
|
||
|
||
self.close()
|
||
try:
|
||
client = mqtt.Client(
|
||
callback_api_version=mqtt.CallbackAPIVersion.VERSION2,
|
||
client_id=f"n7cal_{uuid.uuid4().hex[:12]}",
|
||
protocol=mqtt.MQTTv311,
|
||
)
|
||
except (AttributeError, TypeError):
|
||
client = mqtt.Client(client_id=f"n7cal_{uuid.uuid4().hex[:12]}")
|
||
client.on_connect = self._on_connect
|
||
client.on_disconnect = self._on_disconnect
|
||
client.on_message = self._on_message
|
||
with self._lock:
|
||
self._client = client
|
||
self.on_state("connecting", f"正在连接 {self.host}:{self.port}")
|
||
client.connect_async(self.host, self.port, keepalive=60)
|
||
client.loop_start()
|
||
|
||
def publish_test(self, payload: dict[str, Any] | None = None) -> str:
|
||
with self._lock:
|
||
client = self._client
|
||
if client is None or not client.is_connected():
|
||
raise RuntimeError("MQTT 尚未连接")
|
||
body = payload or sample_payload()
|
||
result = client.publish(self.test_topic, json.dumps(body, ensure_ascii=False), qos=0)
|
||
if result.rc != 0:
|
||
raise RuntimeError(f"MQTT 发送失败，返回码 {result.rc}")
|
||
return self.test_topic
|
||
|
||
def close(self) -> None:
|
||
with self._lock:
|
||
client, self._client = self._client, None
|
||
if client is not None:
|
||
try:
|
||
client.disconnect()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
client.loop_stop()
|
||
except Exception:
|
||
pass
|
||
|
||
def _on_connect(
|
||
self,
|
||
client: Any,
|
||
_userdata: Any,
|
||
_flags: Any,
|
||
reason_code: Any,
|
||
*_: Any,
|
||
) -> None:
|
||
if int(reason_code) != 0:
|
||
self.on_state("error", f"MQTT 连接被拒绝：{reason_code}")
|
||
return
|
||
client.subscribe([(self.topic, 0), (self.test_topic, 0)])
|
||
self.on_state("subscribed", f"已订阅 {self.topic}")
|
||
|
||
def _on_disconnect(self, _client: Any, _userdata: Any, *args: Any) -> None:
|
||
reason = args[-2] if len(args) >= 2 else args[-1] if args else ""
|
||
self.on_state("disconnected", f"MQTT 已断开{f'：{reason}' if reason else ''}")
|
||
|
||
def _on_message(self, _client: Any, _userdata: Any, message: Any) -> None:
|
||
try:
|
||
payload = message.payload.decode("utf-8")
|
||
except UnicodeDecodeError:
|
||
self.on_state("error", "MQTT 收到非 UTF-8 数据")
|
||
return
|
||
self.on_payload(message.topic, payload)
|
||
|
||
|
||
def sample_payload(frame: int = 1) -> dict[str, Any]:
|
||
return {
|
||
"type": "seatpose",
|
||
"frame": frame,
|
||
"timestamp_us": frame * 33_333,
|
||
"area_poses": [
|
||
{
|
||
"area": "Seat_cushion",
|
||
"axis_angle": {"valid": True, "degree": 0.0},
|
||
"current_frame_to_reference": {"valid": True, "T": [0.12, 0.04, 0.01]},
|
||
},
|
||
{
|
||
"area": "Seat_backrest",
|
||
"axis_angle": {"valid": True, "degree": 34.0},
|
||
"current_frame_to_reference": {"valid": True, "T": [0.0, 0.0, 0.0]},
|
||
},
|
||
{
|
||
"area": "Seat_leg_rest",
|
||
"axis_angle": {"valid": True, "degree": 22.0},
|
||
"current_frame_to_reference": {"valid": True, "T": [0.0, 0.0, 0.0]},
|
||
},
|
||
],
|
||
}
|