from dataclasses import dataclass import numpy as np from pygame import Vector2 def checkDimensions(self, row:int, column:int): if self.data.shape != (row, column): raise ValueError(f"{type(self).__name__} must be of shape ({row}, {column}), got {self.data.shape}") # Dataclasses @dataclass class PosArray: data: np.ndarray def __post_init__(self): if not isinstance(self.data, np.ndarray): self.data = np.array(self.data, dtype=float) checkDimensions(self, 6, 3) # behave like numpy def __array__(self): return self.data def __getitem__(self, key): return self.data[key] def __iter__(self): return iter(self.data) def __repr__(self): return f"PosArray(\n{self.data}\n)" def copy(self): return PosArray(self.data.copy()) @dataclass class DegArray: data: np.ndarray def __post_init__(self): if not isinstance(self.data, np.ndarray): self.data = np.array(self.data) checkDimensions(self, 6, 3) self.data = self.data.astype(int) def to_rad(self) -> "RadArray": return RadArray(np.radians(self.data)) # behave like numpy def __array__(self): return self.data def __getitem__(self, key): return self.data[key] def __iter__(self): return iter(self.data) def __repr__(self): return f"PosArray(\n{self.data}\n)" def copy(self): return DegArray(self.data.copy()) @dataclass class RadArray: data: np.ndarray def __post_init__(self): if not isinstance(self.data, np.ndarray): self.data = np.array(self.data, dtype=float) checkDimensions(self, 6, 3) def to_deg(self) -> "DegArray": return DegArray(np.degrees(self.data)) # behave like numpy def __array__(self): return self.data def __getitem__(self, key): return self.data[key] def __iter__(self): return iter(self.data) def __repr__(self): return f"PosArray(\n{self.data}\n)" def copy(self): return RadArray(self.data.copy()) @dataclass class RobotCommand: legs: DegArray state: int = 0 text: str = "" @dataclass class ControlIntent: move_x: float = 0.0 move_y: float = 0.0 turn: float = 0.0 emote: str | None = None walk: bool = False quit: bool = False