feat/refactor: PlayMotion skill.

jws-1 · Mar 11, 2024 · 85b3b05 · 85b3b05
1 parent bbf964e
commit 85b3b05
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 33 deletions.
diff --git a/skills/src/lasr_skills/__init__.py b/skills/src/lasr_skills/__init__.py
@@ -12,3 +12,4 @@
 from .listen_for import ListenFor
 from .receive_object import ReceiveObject
 from .handover_object import HandoverObject
+from .play_motion import PlayMotion
diff --git a/skills/src/lasr_skills/handover_object.py b/skills/src/lasr_skills/handover_object.py
@@ -1,10 +1,9 @@
 import smach
 import smach_ros
 
-from play_motion_msgs.msg import PlayMotionAction, PlayMotionGoal
 from std_srvs.srv import Empty
 
-from lasr_skills import Say, ListenFor
+from lasr_skills import Say, ListenFor, PlayMotion
 
 
 class HandoverObject(smach.StateMachine):
@@ -22,13 +21,7 @@ def __init__(self):
             )
             smach.StateMachine.add(
                 "HANDOVER_OBJECT",
-                smach_ros.SimpleActionState(
-                    "play_motion",
-                    PlayMotionAction,
-                    goal=PlayMotionGoal(
-                        motion_name="handover_object", skip_planning=False
-                    ),
-                ),
+                PlayMotion(motion_name="handover_object"),
                 transitions={"succeeded": "succeeded", "aborted": "failed"},
             )
             smach.StateMachine.add(
@@ -50,19 +43,11 @@ def __init__(self):
             )
             smach.StateMachine.add(
                 "OPEN_GRIPPER",
-                smach_ros.SimpleActionState(
-                    "play_motion",
-                    PlayMotionAction,
-                    goal=PlayMotionGoal(motion_name="open_gripper", skip_planning=True),
-                ),
+                PlayMotion(motion_name="open_gripper"),
                 transitions={"succeeded": "succeeded", "aborted": "failed"},
             )
             smach.StateMachine.add(
                 "FOLD_ARM",
-                smach_ros.SimpleActionState(
-                    "play_motion",
-                    PlayMotionAction,
-                    goal=PlayMotionGoal(motion_name="fold_arm", skip_planning=False),
-                ),
+                PlayMotion(motion_name="fold_arm"),
                 transitions={"succeeded": "succeeded", "aborted": "failed"},
             )
diff --git a/skills/src/lasr_skills/play_motion.py b/skills/src/lasr_skills/play_motion.py
@@ -0,0 +1,24 @@
+import smach_ros
+
+from play_motion_msgs.msg import PlayMotionAction, PlayMotionGoal
+
+from typing import Union
+
+
+class PlayMotion(smach_ros.SimpleActionState):
+    def __init__(self, motion_name: Union[str, None] = None):
+        if motion_name is not None:
+            super(PlayMotion, self).__init__(
+                "play_motion",
+                PlayMotionAction,
+                goal=PlayMotionGoal(motion_name=motion_name, skip_planning=False),
+            )
+        else:
+            super(PlayMotion, self).__init__(
+                "play_motion",
+                PlayMotionAction,
+                goal_cb=lambda ud, _: PlayMotionGoal(
+                    motion_name=ud.motion_name, skip_planning=False
+                ),
+                input_keys=["motion_name"],
+            )
diff --git a/skills/src/lasr_skills/receive_object.py b/skills/src/lasr_skills/receive_object.py
@@ -1,10 +1,9 @@
 import smach
 import smach_ros
 
-from play_motion_msgs.msg import PlayMotionAction, PlayMotionGoal
 from std_srvs.srv import Empty
 
-from lasr_skills import Say, ListenFor
+from lasr_skills import Say, ListenFor, PlayMotion
 
 
 class ReceiveObject(smach.StateMachine):
@@ -22,13 +21,7 @@ def __init__(self):
             )
             smach.StateMachine.add(
                 "RECEIVE_OBJECT",
-                smach_ros.SimpleActionState(
-                    "play_motion",
-                    PlayMotionAction,
-                    goal=PlayMotionGoal(
-                        motion_name="receive_object", skip_planning=False
-                    ),
-                ),
+                PlayMotion(motion_name="receive_object"),
                 transitions={"succeeded": "succeeded", "aborted": "failed"},
             )
             smach.StateMachine.add(
@@ -57,10 +50,6 @@ def __init__(self):
             )
             smach.StateMachine.add(
                 "FOLD_ARM",
-                smach_ros.SimpleActionState(
-                    "play_motion",
-                    PlayMotionAction,
-                    goal=PlayMotionGoal(motion_name="fold_arm", skip_planning=False),
-                ),
+                PlayMotion(motion_name="fold_arm"),
                 transitions={"succeeded": "succeeded", "aborted": "failed"},
             )