From 85b3b05eae4f1e454274f38f365c21f648d0a188 Mon Sep 17 00:00:00 2001
From: Jared Swift <j.w.swift@outlook.com>
Date: Mon, 11 Mar 2024 17:13:46 +0000
Subject: [PATCH] feat/refactor: PlayMotion skill.

---
 skills/src/lasr_skills/__init__.py        |  1 +
 skills/src/lasr_skills/handover_object.py | 23 ++++------------------
 skills/src/lasr_skills/play_motion.py     | 24 +++++++++++++++++++++++
 skills/src/lasr_skills/receive_object.py  | 17 +++-------------
 4 files changed, 32 insertions(+), 33 deletions(-)
 create mode 100644 skills/src/lasr_skills/play_motion.py

diff --git a/skills/src/lasr_skills/__init__.py b/skills/src/lasr_skills/__init__.py
index fc4fcbde9..eadd4d66d 100755
--- 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
index fafd862f3..e3243c3cd 100644
--- 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
new file mode 100644
index 000000000..525e84117
--- /dev/null
+++ 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
index fc3a76479..e4a04be1f 100644
--- 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"},
             )