Improve SRC simulation manager with type_hints

scripts/surgical_robotics_challenge/psm_arm.py

@@ -94,10 +94,10 @@ def __init__(self, simulation_manager:SimulationManager, name, add_joint_errors=
         self.palm_joint_IK = self.simulation_manager.get_obj_handle(name + '_palm_joint_ik')
         self.target_FK = self.simulation_manager.get_obj_handle(name + '_target_fk')
 
-        self.left_finger_ghost = self.simulation_manager._client.get_obj_handle(name + '/left_finger_ghost')
-        self.right_finger_ghost = self.simulation_manager._client.get_obj_handle(name + '/right_finger_ghost')
+        self.left_finger_ghost = self.simulation_manager.get_simulation_ghost(name + '/left_finger_ghost', required=True)
+        self.right_finger_ghost = self.simulation_manager.get_simulation_ghost(name + '/right_finger_ghost', required=True)
         self.actuators = []
-        self.actuators.append(self.simulation_manager._client.get_obj_handle(name + '/Actuator0'))
+        self.actuators.append(self.simulation_manager.get_simulation_actuator(name + '/Actuator0', required=True))
         time.sleep(0.5)
         self.grasped = [False, False, False]
         self.graspable_objs_prefix = ["Needle", "Thread", "Puzzle"]

scripts/surgical_robotics_challenge/simulation_manager.py

@@ -1,11 +1,16 @@
+from typing import Union
 from ambf_client import Client
+from ambf_base_object import BaseObject
+from ambf_actuator import Actuator
+from ambf_rigid_body import RigidBody
+from ambf_ghost_object import GhostObject
 import surgical_robotics_challenge.units_conversion as units_conversion
 
 
 class SimulationObject:
-    def __init__(self, ambf_object):
+    def __init__(self, ambf_object: RigidBody):
         self._object = ambf_object
-        self._joint_type = None # To distinguish between revolute and prismatic joints
+        self._joint_type = None  # To distinguish between revolute and prismatic joints
 
     def set_joint_types(self, joint_types):
         self._joint_type = joint_types
@@ -18,8 +23,8 @@ def get_rotation(self):
 
     def get_pose(self):
         return units_conversion.get_pose(self._object)
-    
-    def get_ros_name(self)->str:
+
+    def get_ros_name(self) -> str:
         return self._object._name
 
     def set_pos(self, pos):
@@ -63,16 +68,52 @@ def __init__(self, name):
         self._client = Client(name)
         self._client.connect()
 
-    def get_obj_handle(self, name, required: bool= False)->SimulationObject:
+    def _get_obj_handle(self, name, required: bool = False) -> Union[BaseObject, None]:
         ambf_object = self._client.get_obj_handle(name)
-        
+
         if required and ambf_object is None:
-            raise RuntimeError(f"SimulationObject {name} is required not found in the simulation")
+            raise RuntimeError(
+                f"Object {name} is required but was not found in the simulation"
+            )
+
+        return ambf_object
+
+    def get_obj_handle(self, name, required: bool = False) -> SimulationObject:
+        ambf_object = self._get_obj_handle(name, required)
 
-        if ambf_object:
+        if isinstance(ambf_object, RigidBody):
             return SimulationObject(ambf_object)
+        elif ambf_object is None and not required:
+            return None
         else:
+            raise RuntimeError(
+                f"Object {name} should be a RigidBody but is instead a {type(ambf_object)}"
+            )
+
+    def get_simulation_actuator(self, name, required: bool = False) -> Actuator:
+        ambf_object = self._get_obj_handle(name, required)
+
+        if isinstance(ambf_object, Actuator):
+            return ambf_object 
+        elif ambf_object is None and not required:
             return None
+        else:
+            raise RuntimeError(
+                f"Object {name} should be a RigidBody but is instead a {type(ambf_object)}"
+            )
+
+
+    def get_simulation_ghost(self, name, required: bool = False) -> GhostObject:
+        ambf_object = self._get_obj_handle(name, required)
+
+        if isinstance(ambf_object, GhostObject):
+            return ambf_object 
+        elif ambf_object is None and not required:
+            return None
+        else:
+            raise RuntimeError(
+                f"Object {name} should be a GhostObject but is instead a {type(ambf_object)}"
+            )
 
     def get_world_handle(self):
         return self._client.get_world_handle()