Suffix asset names when attaching specs.

PiperOrigin-RevId: 794717786
Change-Id: I02abe06ca3c43138e91e61f4acf87a26b3faf7fc

Author: quagla

python/mujoco/specs.cc

@@ -176,6 +176,22 @@ py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) {
   return list;  // list of pointers, so they can be copied
 }
 
+static std::string addSuffixBeforeExtension(const std::string& original_path,
+                                            const std::string& suffix_to_add) {
+  // Find the position of the last dot
+  size_t dot_pos = original_path.rfind('.');
+
+  // Check if a dot was found
+  if (dot_pos != std::string::npos) {
+    std::string new_path = original_path;
+    new_path.insert(dot_pos, suffix_to_add);
+    return new_path;
+  } else {
+    // No extension found, just append
+    return original_path + suffix_to_add;
+  }
+}
+
 PYBIND11_MODULE(_specs, m) {
   auto structs_m = py::module::import("mujoco._structs");
   py::function mjmodel_from_raw_ptr =
@@ -471,13 +487,35 @@ PYBIND11_MODULE(_specs, m) {
             throw pybind11::value_error(mjs_getError(self.ptr));
           }
         }
+        // append the assets adding the suffix to their path
+        std::string suff(s);
         for (const auto& asset : child.assets) {
-          if (self.assets.contains(asset.first) && !self.override_assets) {
+          std::string asset_name =
+              addSuffixBeforeExtension(asset.first.cast<std::string>(), suff);
+          if (self.assets.contains(asset_name) && !self.override_assets) {
             throw pybind11::value_error("Asset " +
                                         asset.first.cast<std::string>() +
                                         " already exists in parent spec.");
           }
-          self.assets[asset.first] = asset.second;
+          self.assets[py::str(asset_name)] = asset.second;
+        }
+        raw::MjsElement* mesh = mjs_firstElement(child.ptr, mjOBJ_MESH);
+        while (mesh) {
+          std::string file = mjs_getString(mjs_asMesh(mesh)->file);
+          if (!file.empty()) {
+            std::string mesh_file = addSuffixBeforeExtension(file, suff);
+            mjs_setString(mjs_asMesh(mesh)->file, mesh_file.c_str());
+          }
+          mesh = mjs_nextElement(child.ptr, mesh);
+        }
+        raw::MjsElement* tex = mjs_firstElement(child.ptr, mjOBJ_TEXTURE);
+        while (tex) {
+          std::string file = mjs_getString(mjs_asTexture(tex)->file);
+          if (!file.empty()) {
+            std::string tex_file = addSuffixBeforeExtension(file, suff);
+            mjs_setString(mjs_asTexture(tex)->file, tex_file.c_str());
+          }
+          tex = mjs_nextElement(child.ptr, tex);
         }
         child.parent = &self;
         return mjs_asFrame(attached_frame);

python/mujoco/specs_test.py

@@ -1069,7 +1069,7 @@ def test_attach_to_site(self):
     child2 = mujoco.MjSpec()
     child2.assets = {'cube2.obj': 'cube2_content'}
     body2 = child2.worldbody.add_body(name='body')
-    self.assertIsNotNone(parent.attach(child2, site=site, prefix='child2-'))
+    self.assertIsNotNone(parent.attach(child2, site=site, suffix='-child2'))
     self.assertIsNotNone(child2.worldbody)
     self.assertEqual(child2.parent, parent)
     body2.pos = [-1, -1, -1]
@@ -1081,13 +1081,13 @@ def test_attach_to_site(self):
     np.testing.assert_array_equal(model2.body_quat[1], [0, 0, 0, 1])
     np.testing.assert_array_equal(model2.body_quat[2], [0, 0, 0, 1])
     self.assertEqual(parent.assets['cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2.obj'], 'cube2_content')
+    self.assertEqual(parent.assets['cube2-child2.obj'], 'cube2_content')
 
     # Attach another spec to site (referenced by name) and compile again.
     child3 = mujoco.MjSpec()
     child3.assets = {'cube3.obj': 'cube3_content'}
     body3 = child3.worldbody.add_body(name='body')
-    self.assertIsNotNone(parent.attach(child3, site='site', prefix='child3-'))
+    self.assertIsNotNone(parent.attach(child3, site='site', suffix='-child3'))
     self.assertIsNotNone(child3.worldbody)
     self.assertEqual(child3.parent, parent)
     body3.pos = [-2, -2, -2]
@@ -1101,8 +1101,8 @@ def test_attach_to_site(self):
     np.testing.assert_array_equal(model3.body_quat[2], [0, 0, 0, 1])
     np.testing.assert_array_equal(model3.body_quat[3], [0, 0, 0, 1])
     self.assertEqual(parent.assets['cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2.obj'], 'cube2_content')
-    self.assertEqual(parent.assets['cube3.obj'], 'cube3_content')
+    self.assertEqual(parent.assets['cube2-child2.obj'], 'cube2_content')
+    self.assertEqual(parent.assets['cube3-child3.obj'], 'cube3_content')
 
     # Fail to attach to a site that does not exist.
     child4 = mujoco.MjSpec()
@@ -1150,7 +1150,7 @@ def test_attach_to_frame(self):
     child2.assets = {'cube2.obj': 'cube2_content'}
     body2 = child2.worldbody.add_body(name='body')
     body2.set_frame(child2.worldbody.add_frame(pos=[-1, -1, 1]))
-    self.assertIsNotNone(parent.attach(child2, frame=frame, prefix='child-'))
+    self.assertIsNotNone(parent.attach(child2, frame=frame, suffix='-child'))
     self.assertIsNotNone(child2.worldbody)
     self.assertEqual(child2.parent, parent)
     body2.pos = [-1, -1, -1]
@@ -1162,14 +1162,14 @@ def test_attach_to_frame(self):
     np.testing.assert_array_equal(model2.body_quat[1], [0, 0, 0, 1])
     np.testing.assert_array_equal(model2.body_quat[2], [0, 0, 0, 1])
     self.assertEqual(parent.assets['cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2.obj'], 'cube2_content')
+    self.assertEqual(parent.assets['cube2-child.obj'], 'cube2_content')
 
     # Attach another spec to frame (referenced by name) and compile again.
     child3 = mujoco.MjSpec()
     child3.assets = {'cube2.obj': 'new_cube2_content'}
     body3 = child3.worldbody.add_body(name='body')
     body3.set_frame(child3.worldbody.add_frame(pos=[-1, -1, 1]))
-    self.assertIsNotNone(parent.attach(child3, frame='frame', prefix='child3-'))
+    self.assertIsNotNone(parent.attach(child3, frame='frame', suffix='-child3'))
     self.assertIsNotNone(child3.worldbody)
     self.assertEqual(child3.parent, parent)
     body3.pos = [-2, -2, -2]
@@ -1183,7 +1183,7 @@ def test_attach_to_frame(self):
     np.testing.assert_array_equal(model3.body_quat[2], [0, 0, 0, 1])
     np.testing.assert_array_equal(model3.body_quat[3], [0, 0, 0, 1])
     self.assertEqual(parent.assets['cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2.obj'], 'new_cube2_content')
+    self.assertEqual(parent.assets['cube2-child3.obj'], 'new_cube2_content')
 
     # Fail to attach to a frame that does not exist.
     child4 = mujoco.MjSpec()