From f8a4f4eb2311cca7788ad754453aefdfaf3db564 Mon Sep 17 00:00:00 2001 From: Benjamin Bolte Date: Wed, 16 Oct 2024 10:46:48 -0700 Subject: [PATCH] copy denys' code over --- openlch/requirements.txt | 2 + openlch/websocket.py | 233 +++++++++++++++++++++++++++++++++++++++ tests/test_dummy.py | 12 -- tests/test_websocket.py | 5 + 4 files changed, 240 insertions(+), 12 deletions(-) create mode 100644 openlch/websocket.py delete mode 100644 tests/test_dummy.py create mode 100644 tests/test_websocket.py diff --git a/openlch/requirements.txt b/openlch/requirements.txt index 2ee1e7b..591bd36 100644 --- a/openlch/requirements.txt +++ b/openlch/requirements.txt @@ -1,2 +1,4 @@ # requirements.txt +websockets +pydantic diff --git a/openlch/websocket.py b/openlch/websocket.py new file mode 100644 index 0000000..04ebf43 --- /dev/null +++ b/openlch/websocket.py @@ -0,0 +1,233 @@ +"""Defines the ServoClient class for controlling the OpenLCH servo motors.""" + +import asyncio +import json +from enum import Enum +from typing import Any + +import websockets +from pydantic import BaseModel + + +class ServoMode(Enum): + Position = 0 + ConstantSpeed = 1 + PWMOpenLoop = 2 + StepServo = 3 + + +class ServoDirection(Enum): + Clockwise = 0 + Counterclockwise = 1 + + +class MemoryLockState(Enum): + Unlocked = 0 + Locked = 1 + + +class TorqueMode(Enum): + Disabled = 0 + Enabled = 1 + Stiff = 2 + + +class ServoData(BaseModel): + acceleration: int + async_write_flag: int + current_current: int + current_load: int + current_location: int + current_speed: int + current_temperature: int + current_voltage: int + lock_mark: int + mobile_sign: int + reserved1: int + reserved2: int + running_speed: int + running_time: int + servo_status: int + target_location: int + torque_limit: int + torque_switch: int + + def __repr__(self) -> str: + attrs = [ + f"current_location={self.current_location}", + f"current_speed={self.current_speed}", + f"current_temperature={self.current_temperature}", + ] + return f"ServoData({', '.join(attrs)})" + + +class ContinuousReadResponse(BaseModel): + servo_data: list[ServoData] + task_run_count: int + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ContinuousReadResponse": + return cls(servo_data=[ServoData(**servo) for servo in data["servo"]], task_run_count=data["task_run_count"]) + + def __repr__(self) -> str: + attrs = [ + f"servos={len(self.servo_data)}", + f"task_run_count={self.task_run_count}", + ] + return f"ContinuousReadResponse({', '.join(attrs)})" + + +class WriteMultipleCommand(BaseModel): + ids: list[int] + positions: list[int] + times: list[int] + speeds: list[int] + only_write_positions: int + + def __repr__(self) -> str: + attrs = [ + f"ids={self.ids}", + f"positions={self.positions}", + f"times={self.times}", + f"speeds={self.speeds}", + f"only_write_positions={self.only_write_positions}", + ] + return f"WriteMultipleCommand({', '.join(attrs)})" + + +class ServoClient: + def __init__(self, uri: str = "ws://192.168.42.1:8080") -> None: + self.uri = uri + self.websocket = None + + async def connect(self) -> None: + if self.websocket is None or self.websocket.closed: + self.websocket = await websockets.connect(self.uri) + + async def disconnect(self) -> None: + if self.websocket and not self.websocket.closed: + await self.websocket.close() + + async def _send_command(self, command: str, params: dict[str, Any] | None) -> dict[str, Any] | None: + try: + await self.connect() + message = json.dumps({"command": command, "params": params}) + await self.websocket.send(message) + response = await self.websocket.recv() + parsed_response = json.loads(response) + + if not parsed_response["success"]: + error_msg = parsed_response.get("error", "Unknown error") + print(f"Error in command {command}: {error_msg}") + return None + + return parsed_response.get("data") + except websockets.exceptions.WebSocketException as e: + print(f"WebSocket error: {e}") + await self.disconnect() + return None + except json.JSONDecodeError as e: + print(f"JSON decoding error: {e}") + return None + except Exception as e: + print(f"Unexpected error: {e}") + await self.disconnect() + return None + + async def move_servo(self, id: int, position: int, time: int, speed: int) -> dict[str, Any] | None: + params = {"id": id, "position": position, "time": time, "speed": speed} + return await self._send_command("Move", params) + + async def set_mode(self, id: int, mode: ServoMode) -> dict[str, Any] | None: + params = {"id": id, "mode": mode.value} + return await self._send_command("SetMode", params) + + async def set_speed(self, id: int, speed: int, direction: ServoDirection) -> dict[str, Any] | None: + params = {"id": id, "speed": speed, "direction": direction.value} + return await self._send_command("SetSpeed", params) + + async def read_info(self, id: int) -> dict[str, Any] | None: + params = {"id": id} + return await self._send_command("ReadInfo", params) + + async def read_pid(self, id: int) -> dict[str, Any] | None: + params = {"id": id} + return await self._send_command("ReadPID", params) + + async def set_pid(self, id: int, p: int, i: int, d: int) -> dict[str, Any] | None: + params = {"id": id, "p": p, "i": i, "d": d} + return await self._send_command("SetPID", params) + + async def set_memory_lock(self, id: int, state: MemoryLockState) -> dict[str, Any] | None: + params = {"id": id, "state": state.value} + return await self._send_command("SetMemoryLock", params) + + async def read_angle_limits(self, id: int) -> dict[str, Any] | None: + params = {"id": id} + return await self._send_command("ReadAngleLimits", params) + + async def set_torque_mode(self, id: int, mode: TorqueMode) -> dict[str, Any] | None: + params = {"id": id, "mode": mode.value} + return await self._send_command("SetTorqueMode", params) + + async def scan(self, id: int) -> dict[str, Any] | None: + params = {"id": id} + return await self._send_command("Scan", params) + + async def write_multiple(self, command: WriteMultipleCommand) -> dict[str, Any] | None: + params = {"cmd": command.dict()} + return await self._send_command("WriteMultiple", params) + + async def read_continuous(self) -> ContinuousReadResponse | None: + data = await self._send_command("ReadContinuous", None) + if data: + return ContinuousReadResponse(data) + return None + + +async def main() -> None: + client = ServoClient() + + odd = True + move_timer = 0 + target_interval = 1 / 50 # 50Hz + try: + while True: + loop_start = asyncio.get_event_loop().time() + + # Read continuous data + response = await client.read_continuous() + if response: + print("Read continuous response:", response.servo_data[8]) # Print data for servo 9 (index 8) + + # Move servo once per second + if move_timer >= 50: # 50 iterations at 50Hz = 1 second + await client.write_multiple( + WriteMultipleCommand( + ids=list(range(1, 17)), + positions=[3500] * 16 if odd else [3000] * 16, + times=[0] * 16, + speeds=[0] * 16, + only_write_positions=1, + ) + ) + odd = not odd + move_timer = 0 + else: + move_timer += 1 + + # Calculate remaining time and sleep if possible + elapsed_time = asyncio.get_event_loop().time() - loop_start + remaining_time = target_interval - elapsed_time + + if remaining_time > 0: + await asyncio.sleep(remaining_time) + else: + print(f"Warning: Loop running late by {-remaining_time:.6f} seconds") + + finally: + await client.disconnect() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/test_dummy.py b/tests/test_dummy.py deleted file mode 100644 index 284dabe..0000000 --- a/tests/test_dummy.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Defines a dummy test.""" - -import pytest - - -def test_dummy() -> None: - assert True - - -@pytest.mark.slow -def test_slow() -> None: - assert True diff --git a/tests/test_websocket.py b/tests/test_websocket.py new file mode 100644 index 0000000..cc40a96 --- /dev/null +++ b/tests/test_websocket.py @@ -0,0 +1,5 @@ +"""Tests the websocket client (currently not implemented).""" + + +def test_websocket_client() -> None: + assert False