Merge branch 'xarm-rt-driver' of https://github.com/dimensionalOS/dimos into xarm-rt-driver

Author: mustafab0

dimos/hardware/manipulators/xarm/components/__init__.py

@@ -1,15 +1,15 @@
 """Component classes for XArmDriver."""
 
+from .gripper_control import GripperControlComponent
+from .kinematics import KinematicsComponent
 from .motion_control import MotionControlComponent
 from .state_queries import StateQueryComponent
 from .system_control import SystemControlComponent
-from .kinematics import KinematicsComponent
-from .gripper_control import GripperControlComponent
 
 __all__ = [
+    "GripperControlComponent",
+    "KinematicsComponent",
     "MotionControlComponent",
     "StateQueryComponent",
     "SystemControlComponent",
-    "KinematicsComponent",
-    "GripperControlComponent",
 ]

dimos/hardware/manipulators/xarm/interactive_control.py

@@ -27,30 +27,30 @@
     venv/bin/python dimos/hardware/manipulators/xarm/interactive_control.py
 """
 
+import math
 import os
 import time
-import math
 
 from dimos import core
-from dimos.hardware.manipulators.xarm.xarm_driver import XArmDriver
 from dimos.hardware.manipulators.xarm.sample_trajectory_generator import (
     SampleTrajectoryGenerator,
 )
-from dimos.msgs.sensor_msgs import JointState, RobotState, JointCommand
+from dimos.hardware.manipulators.xarm.xarm_driver import XArmDriver
+from dimos.msgs.sensor_msgs import JointCommand, JointState, RobotState
 from dimos.utils.logging_config import setup_logger
 
 logger = setup_logger(__file__)
 
 
-def print_banner():
+def print_banner() -> None:
     """Print welcome banner."""
     print("\n" + "=" * 80)
     print("  xArm Interactive Control")
     print("  Real-time joint control via terminal UI")
     print("=" * 80)
 
 
-def print_current_state(traj_gen):
+def print_current_state(traj_gen) -> None:
     """Display current joint positions."""
     state = traj_gen.get_current_state()
 
@@ -99,7 +99,7 @@ def get_control_mode():
             return None
 
 
-def get_joint_selection(num_joints):
+def get_joint_selection(num_joints: int):
     """Get joint selection from user."""
     while True:
         try:
@@ -240,14 +240,14 @@ def confirm_motion_velocity(joint_index, velocity_deg_s, duration):
         f"  Velocity: {velocity_deg_s:+.2f}°/s ({'clockwise' if velocity_deg_s < 0 else 'counterclockwise'})"
     )
     print(f"  Duration: {duration:.2f}s (with ramp up/down)")
-    print(f"  Profile: 20% ramp up, 60% constant, 20% ramp down")
+    print("  Profile: 20% ramp up, 60% constant, 20% ramp down")
     print("=" * 80)
 
     confirm = input("\nExecute this motion? (y/n): ").strip().lower()
     return confirm == "y"
 
 
-def wait_for_trajectory_completion(traj_gen, duration):
+def wait_for_trajectory_completion(traj_gen, duration) -> None:
     """Wait for trajectory to complete and show progress."""
     print("\n⚙ Executing motion...")
 
@@ -272,7 +272,7 @@ def wait_for_trajectory_completion(traj_gen, duration):
     print("✓ Motion complete!")
 
 
-def interactive_control_loop(xarm, traj_gen, num_joints):
+def interactive_control_loop(xarm, traj_gen, num_joints: int) -> None:
     """Main interactive control loop."""
     print_banner()
 
@@ -424,7 +424,7 @@ def interactive_control_loop(xarm, traj_gen, num_joints):
     print("=" * 80)
 
 
-def main():
+def main() -> None:
     """Run interactive xArm control."""
     import argparse
 

dimos/hardware/manipulators/xarm/sample_trajectory_generator.py

@@ -39,14 +39,13 @@
     traj_gen.start()
 """
 
-import time
-import threading
-import math
-from typing import List, Optional
 from dataclasses import dataclass
+import math
+import threading
+import time
 
-from dimos.core import Module, ModuleConfig, In, Out, rpc
-from dimos.msgs.sensor_msgs import JointState, RobotState, JointCommand
+from dimos.core import In, Module, ModuleConfig, Out, rpc
+from dimos.msgs.sensor_msgs import JointCommand, JointState, RobotState
 from dimos.utils.logging_config import setup_logger
 
 logger = setup_logger(__file__)
@@ -85,18 +84,18 @@ class SampleTrajectoryGenerator(Module):
     joint_position_command: Out[JointCommand] = None  # Position commands (radians)
     joint_velocity_command: Out[JointCommand] = None  # Velocity commands (rad/s)
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         super().__init__(**kwargs)
 
         # State tracking
-        self._current_joint_state: Optional[JointState] = None
-        self._current_robot_state: Optional[RobotState] = None
+        self._current_joint_state: JointState | None = None
+        self._current_robot_state: RobotState | None = None
         self._state_lock = threading.Lock()
 
         # Control thread
         self._running = False
         self._stop_event = threading.Event()
-        self._control_thread: Optional[threading.Thread] = None
+        self._control_thread: threading.Thread | None = None
 
         # Publishing enabled flag
         self._publishing_enabled = self.config.enable_on_start
@@ -121,7 +120,7 @@ def __init__(self, **kwargs):
         )
 
     @rpc
-    def start(self):
+    def start(self) -> None:
         """Start the trajectory generator."""
         super().start()
 
@@ -144,7 +143,7 @@ def start(self):
         logger.info("Trajectory generator started")
 
     @rpc
-    def stop(self):
+    def stop(self) -> None:
         """Stop the trajectory generator."""
         logger.info("Stopping trajectory generator...")
 
@@ -159,19 +158,19 @@ def stop(self):
         logger.info("Trajectory generator stopped")
 
     @rpc
-    def enable_publishing(self):
+    def enable_publishing(self) -> None:
         """Enable command publishing."""
         self._publishing_enabled = True
         logger.info("Command publishing enabled")
 
     @rpc
-    def disable_publishing(self):
+    def disable_publishing(self) -> None:
         """Disable command publishing."""
         self._publishing_enabled = False
         logger.info("Command publishing disabled")
 
     @rpc
-    def set_velocity_mode(self, enabled: bool):
+    def set_velocity_mode(self, enabled: bool) -> None:
         """
         Set velocity mode flag.
 
@@ -313,7 +312,7 @@ def move_joint_velocity(self, joint_index: int, velocity_deg_s: float, duration:
     # Private Methods: Callbacks
     # =========================================================================
 
-    def _on_joint_state(self, msg: JointState):
+    def _on_joint_state(self, msg: JointState) -> None:
         """Callback for receiving joint state updates."""
         with self._state_lock:
             # Log first message with all joints
@@ -329,7 +328,7 @@ def _on_joint_state(self, msg: JointState):
                 )
             self._current_joint_state = msg
 
-    def _on_robot_state(self, msg: RobotState):
+    def _on_robot_state(self, msg: RobotState) -> None:
         """Callback for receiving robot state updates."""
         with self._state_lock:
             # Log first message or when state/error changes
@@ -360,7 +359,7 @@ def _on_robot_state(self, msg: RobotState):
     # Private Methods: Control Loop
     # =========================================================================
 
-    def _start_control_loop(self):
+    def _start_control_loop(self) -> None:
         """Start the control loop thread."""
         logger.info(f"Starting control loop at {self.config.publish_rate}Hz")
 
@@ -372,7 +371,7 @@ def _start_control_loop(self):
         )
         self._control_thread.start()
 
-    def _control_loop(self):
+    def _control_loop(self) -> None:
         """
         Control loop for publishing commands.
 
@@ -436,7 +435,7 @@ def _control_loop(self):
 
         logger.info("Control loop stopped")
 
-    def _generate_command(self) -> Optional[List[float]]:
+    def _generate_command(self) -> list[float] | None:
         """
         Generate command for the robot.
 
@@ -514,7 +513,7 @@ def _generate_command(self) -> Optional[List[float]]:
                 # Position mode with no trajectory: hold current position (safe)
                 return list(self._current_joint_state.position)
 
-    def _publish_position_command(self, command: List[float]):
+    def _publish_position_command(self, command: list[float]) -> None:
         """Publish joint position command."""
         if self.joint_position_command._transport or (
             hasattr(self.joint_position_command, "connection")
@@ -539,7 +538,7 @@ def _publish_position_command(self, command: List[float]):
             if self._command_count == 0:
                 logger.warning("joint_position_command transport not configured!")
 
-    def _publish_velocity_command(self, command: List[float]):
+    def _publish_velocity_command(self, command: list[float]) -> None:
         """Publish joint velocity command."""
         if self.joint_velocity_command._transport or (
             hasattr(self.joint_velocity_command, "connection")

dimos/hardware/manipulators/xarm/spec.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Protocol, List, Tuple
+from typing import Protocol
 
 from dimos.core import In, Out
-from dimos.msgs.geometry_msgs import PoseStamped, Twist, WrenchStamped
-from dimos.msgs.nav_msgs import Path
-from dimos.msgs.sensor_msgs import JointState, RobotState, JointCommand
+from dimos.msgs.geometry_msgs import WrenchStamped
+from dimos.msgs.sensor_msgs import JointCommand, JointState, RobotState
 
 
 class ArmDriverSpec(Protocol):
@@ -37,14 +36,14 @@ class ArmDriverSpec(Protocol):
     ft_raw: Out[WrenchStamped]  # Raw force/torque sensor data
 
     # RPC Methods
-    def set_joint_angles(self, angles: List[float]) -> Tuple[int, str]: ...
+    def set_joint_angles(self, angles: list[float]) -> tuple[int, str]: ...
 
-    def set_joint_velocities(self, velocities: List[float]) -> Tuple[int, str]: ...
+    def set_joint_velocities(self, velocities: list[float]) -> tuple[int, str]: ...
 
     def get_joint_state(self) -> JointState: ...
 
     def get_robot_state(self) -> RobotState: ...
 
-    def enable_servo_mode(self) -> Tuple[int, str]: ...
+    def enable_servo_mode(self) -> tuple[int, str]: ...
 
-    def disable_servo_mode(self) -> Tuple[int, str]: ...
+    def disable_servo_mode(self) -> tuple[int, str]: ...

dimos/hardware/manipulators/xarm/test_xarm_driver.py

@@ -38,8 +38,8 @@
 
 from dimos import core
 from dimos.hardware.manipulators.xarm.xarm_driver import XArmDriver
-from dimos.msgs.sensor_msgs import JointState, RobotState
 from dimos.msgs.geometry_msgs import WrenchStamped
+from dimos.msgs.sensor_msgs import JointState, RobotState
 from dimos.utils.logging_config import setup_logger
 
 logger = setup_logger(__file__)
@@ -108,7 +108,7 @@ def driver(dimos_cluster):
 
 
 @pytest.mark.tool
-def test_basic_connection(driver):
+def test_basic_connection(driver) -> None:
     """Test basic connection and startup with dimos deployment."""
     logger.info("=" * 80)
     logger.info("TEST 1: Basic Connection with dimos.deploy()")
@@ -135,7 +135,7 @@ def test_basic_connection(driver):
 
 
 @pytest.mark.tool
-def test_joint_state_reading(driver):
+def test_joint_state_reading(driver) -> None:
     """Test joint state reading via LCM topic subscription."""
     logger.info("=" * 80)
     logger.info("TEST 2: Joint State Reading via LCM Transport")
@@ -149,7 +149,7 @@ def test_joint_state_reading(driver):
     joint_states_received = []
     robot_states_received = []
 
-    def on_joint_state(msg):
+    def on_joint_state(msg) -> None:
         """Callback for receiving joint state messages from LCM."""
         joint_states_received.append(msg)
         if len(joint_states_received) <= 3:
@@ -159,7 +159,7 @@ def on_joint_state(msg):
                 f"(showing first 3 joints)"
             )
 
-    def on_robot_state(msg):
+    def on_robot_state(msg) -> None:
         """Callback for receiving robot state messages from LCM."""
         robot_states_received.append(msg)
         if len(robot_states_received) <= 3:
@@ -227,7 +227,7 @@ def on_robot_state(msg):
 
 
 @pytest.mark.tool
-def test_command_sending(driver):
+def test_command_sending(driver) -> None:
     """Test that command RPC methods are available and functional."""
     logger.info("=" * 80)
     logger.info("TEST 3: Command RPC Methods")
@@ -271,7 +271,7 @@ def test_command_sending(driver):
 
 
 @pytest.mark.tool
-def test_rpc_methods(driver):
+def test_rpc_methods(driver) -> None:
     """Test RPC method calls."""
     logger.info("=" * 80)
     logger.info("TEST 4: RPC Methods")
@@ -304,7 +304,7 @@ def test_rpc_methods(driver):
     logger.info("✓ TEST 4 PASSED\n")
 
 
-def run_tests():
+def run_tests() -> None:
     """Run all tests."""
     logger.info("=" * 80)
     logger.info("XArm Driver Test Suite (Full dimos Deployment)")
@@ -371,7 +371,7 @@ def run_tests():
         logger.error("❌ SOME TESTS FAILED")
 
 
-def run_driver():
+def run_driver() -> None:
     """Start the xArm driver and keep it running."""
     logger.info("=" * 80)
     logger.info("XArm Driver - Starting in continuous mode")
@@ -431,9 +431,8 @@ def run_driver():
         logger.info("✓ Driver stopped")
 
 
-def main():
+def main() -> None:
     """Main entry point."""
-    import sys
     import argparse
 
     # Parse command-line arguments