Merge pull request #95 from iory/eus

Refactor eus model server

ros/kxreus/node_scripts/eus_model_server.py

@@ -2,6 +2,7 @@
 
 import os
 import tempfile
+from contextlib import contextmanager
 
 from filelock import FileLock
 from kxr_models.md5sum_utils import checksum_md5
@@ -15,84 +16,91 @@
 
 class EusModelServer:
     def __init__(self):
+        self.namespace = self._get_clean_namespace()
+        self.kxr_models_path = rospkg.RosPack().get_path("kxr_models")
+        self.previous_md5sum = None
+        self.robot_model = None
+
+    def _get_clean_namespace(self):
         full_namespace = rospy.get_namespace()
         last_slash_pos = full_namespace.rfind("/")
-        self.clean_namespace = (
-            full_namespace[:last_slash_pos] if last_slash_pos != 0 else ""
-        )
+        return full_namespace[:last_slash_pos] if last_slash_pos != 0 else ""
 
-    def load_robot_model(self, urdf_path):
+    @staticmethod
+    def load_robot_model(urdf_content):
+        """Load URDF content into a RobotModel instance."""
         r = RobotModel()
-        with open(urdf_path) as f:
-            with no_mesh_load_mode():
+        with EusModelServer.temp_file() as urdf_file:
+            with open(urdf_file, "w") as f:
+                f.write(urdf_content)
+            with open(urdf_file) as f, no_mesh_load_mode():
                 r.load_urdf_file(f)
         return r
 
-    def run(self):
-        rate = rospy.Rate(1)
+    @staticmethod
+    @contextmanager
+    def temp_file(suffix=""):
+        temp = tempfile.mktemp(suffix=suffix)
+        try:
+            yield temp
+        finally:
+            if os.path.exists(temp):
+                os.remove(temp)
+
+    def get_md5sum(self, urdf_content, joint_group_description):
+        with self.temp_file() as urdf_file:
+            with open(urdf_file, "w") as f:
+                f.write(urdf_content)
+            md5sum = checksum_md5(urdf_file)
+
+        if joint_group_description:
+            with self.temp_file(suffix=".yaml") as yaml_file:
+                with open(yaml_file, "w") as f:
+                    yaml.dump(joint_group_description, f)
+                md5sum += "-" + checksum_md5(yaml_file)
+
+        return md5sum
 
-        previous_md5sum = None
-        robot_model = None
+    def save_eus_model(self, md5sum, urdf_content, joint_group_description):
+        eus_path = os.path.join(
+            self.kxr_models_path, "models", "euslisp", f"{md5sum}.l"
+        )
+        if os.path.exists(eus_path):
+            return eus_path  # Model already exists
+
+        lock = FileLock(eus_path + ".lock", timeout=10)
+        with lock:
+            with self.temp_file() as urdf_file:
+                with open(urdf_file, "w") as f:
+                    f.write(urdf_content)
+                with open(eus_path, "w") as f:
+                    urdf2eus(urdf_file, config_yaml_path=joint_group_description, fp=f)
+            rospy.loginfo(f"Eusmodel is saved to {eus_path}")
+        return eus_path
 
-        rospack = rospkg.RosPack()
-        kxr_models_path = rospack.get_path("kxr_models")
+    def run(self):
+        rate = rospy.Rate(1)
 
         while not rospy.is_shutdown():
             rate.sleep()
-            urdf = rospy.get_param(self.clean_namespace + "/robot_description", None)
-            if urdf is not None:
-                tmp_file = tempfile.mktemp()
-                with open(tmp_file, "w") as f:
-                    f.write(urdf)
+            urdf_content = rospy.get_param(self.namespace + "/robot_description", None)
+            if urdf_content:
                 joint_group_description = rospy.get_param(
-                    self.clean_namespace + "/joint_group_description", None
+                    self.namespace + "/joint_group_description", None
                 )
-                md5sum = checksum_md5(tmp_file)
-                tmp_joint_description_file = None
-                if joint_group_description is not None:
-                    tmp_joint_description_file = tempfile.mktemp(suffix=".yaml")
-                    with open(tmp_joint_description_file, "w") as f:
-                        yaml.dump(joint_group_description, f)
-                    md5sum = md5sum + '-' + checksum_md5(tmp_joint_description_file)
-
-                if previous_md5sum is None or previous_md5sum != md5sum:
-                    robot_model = self.load_robot_model(tmp_file)
-                previous_md5sum = md5sum
-
-                eus_path = os.path.join(
-                    kxr_models_path, "models", "euslisp", f"{md5sum}.l"
+                md5sum = self.get_md5sum(urdf_content, joint_group_description)
+
+                if self.previous_md5sum != md5sum:
+                    self.robot_model = self.load_robot_model(urdf_content)
+                    self.previous_md5sum = md5sum
+
+                self.save_eus_model(
+                    md5sum, urdf_content, joint_group_description
                 )
-                robot_name = robot_model.urdf_robot_model.name
-                if os.path.exists(eus_path):
-                    rospy.set_param(
-                        self.clean_namespace + "/eus_robot_name", robot_name
-                    )
-                    rospy.set_param(self.clean_namespace + "/eusmodel_hash", md5sum)
-                    continue
-
-                lock_path = eus_path + ".lock"
-                lock = FileLock(lock_path, timeout=10)
-                try:
-                    with lock:
-                        with open(eus_path, "w") as f:
-                            urdf2eus(
-                                tmp_file,
-                                config_yaml_path=tmp_joint_description_file,
-                                fp=f,
-                            )
-                        rospy.loginfo(f"Eusmodel is saved to {eus_path}")
-                    os.remove(lock_path)
-                    rospy.set_param(
-                        self.clean_namespace + "/eus_robot_name", robot_name
-                    )
-                    rospy.set_param(self.clean_namespace + "/eusmodel_hash", md5sum)
-                finally:
-                    if os.path.exists(tmp_file):
-                        os.remove(tmp_file)
-                    if tmp_joint_description_file is not None and os.path.exists(
-                        tmp_joint_description_file
-                    ):
-                        os.remove(tmp_joint_description_file)
+                robot_name = self.robot_model.urdf_robot_model.name
+
+                rospy.set_param(self.namespace + "/eus_robot_name", robot_name)
+                rospy.set_param(self.namespace + "/eusmodel_hash", md5sum)
 
 
 if __name__ == "__main__":