Merge branch 'main' into modular-graphnav

bdaiinstitute · Jul 31, 2023 · b16ddc3 · b16ddc3
2 parents 9b0eedf + c85fe4a
commit b16ddc3
Show file tree

Hide file tree

Showing 4 changed files with 760 additions and 48 deletions.
diff --git a/spot_wrapper/spot_arm.py b/spot_wrapper/spot_arm.py
@@ -21,7 +21,7 @@
 from bosdyn.client.time_sync import TimeSyncEndpoint
 from bosdyn.util import seconds_to_duration
 
-from spot_wrapper.wrapper_helpers import RobotState
+from spot_wrapper.wrapper_helpers import RobotState, ClaimAndPowerDecorator
 
 
 class SpotArm:
@@ -34,13 +34,19 @@ def __init__(
         manipulation_api_client: ManipulationApiClient,
         robot_state_client: RobotStateClient,
         max_command_duration: float,
+        claim_and_power_decorator: ClaimAndPowerDecorator,
     ) -> None:
         """
         Constructor for SpotArm class.
         Args:
             robot: Robot object
             logger: Logger object
+            robot_state: Object containing the robot's state as controlled by the wrapper
+            robot_command_client: Command client to use to send commands to the robot
+            manipulation_api_client: Command client to send manipulation commands to the robot
+            robot_state_client: Client to retrieve state of the robot
             max_command_duration: Maximum duration for commands when using the manipulation command method
+            claim_and_power_decorator: Object to use to decorate the functions on this object
         """
         self._robot = robot
         self._logger = logger
@@ -49,6 +55,23 @@ def __init__(
         self._robot_command_client = robot_command_client
         self._manipulation_api_client = manipulation_api_client
         self._robot_state_client = robot_state_client
+        self._claim_and_power_decorator = claim_and_power_decorator
+        self._claim_and_power_decorator.decorate_functions(
+            self,
+            decorated_funcs=[
+                self.ensure_arm_power_and_stand,
+                self.arm_stow,
+                self.arm_unstow,
+                self.arm_carry,
+                self.arm_joint_move,
+                self.force_trajectory,
+                self.gripper_open,
+                self.gripper_close,
+                self.gripper_angle_open,
+                self.hand_pose,
+                self.grasp_3d,
+            ],
+        )
 
     def _manipulation_request(
         self,

diff --git a/spot_wrapper/spot_docking.py b/spot_wrapper/spot_docking.py
@@ -6,7 +6,11 @@
 from bosdyn.client.docking import DockingClient, blocking_dock_robot, blocking_undock
 from bosdyn.client.robot import Robot
 
-from spot_wrapper.wrapper_helpers import RobotState, RobotCommandData
+from spot_wrapper.wrapper_helpers import (
+    RobotState,
+    RobotCommandData,
+    ClaimAndPowerDecorator,
+)
 
 
 class SpotDocking:
@@ -22,13 +26,20 @@ def __init__(
         command_data: RobotCommandData,
         docking_client: DockingClient,
         robot_command_client: robot_command.RobotCommandClient,
+        claim_and_power_decorator: ClaimAndPowerDecorator,
     ) -> None:
         self._robot = robot
         self._logger = logger
         self._command_data = command_data
         self._docking_client: DockingClient = docking_client
         self._robot_command_client = robot_command_client
         self._robot_state = robot_state
+        self._claim_and_power_decorator = claim_and_power_decorator
+        # Decorate the functions so that they take the lease. Dock function needs to power on because it might have
+        # to move the robot, the undock
+        self._claim_and_power_decorator.decorate_functions(
+            self, decorated_funcs=[self.dock, self.undock]
+        )
 
     def dock(self, dock_id: int) -> typing.Tuple[bool, str]:
         """Dock the robot to the docking station with fiducial ID [dock_id]."""