Add prefix to asset paths when appending specs.

PiperOrigin-RevId: 794966246
Change-Id: I4825a11cc3234e6f403340dd997742cea0e2183b

python/mujoco/specs.cc

@@ -192,6 +192,29 @@ static std::string addSuffixBeforeExtension(const std::string& original_path,
   }
 }
 
+static std::string addPrefixToFileName(const std::string& original_path,
+                                       const std::string& prefix_to_add) {
+  // Find the position of the last slash
+  size_t slash_pos = original_path.find_last_of("/\\");
+
+  // Check if a slash was found
+  if (slash_pos != std::string::npos) {
+    std::string new_path = original_path;
+    new_path.insert(slash_pos + 1, prefix_to_add);
+    return new_path;
+  } else {
+    // No slash found, just prepend
+    return prefix_to_add + original_path;
+  }
+}
+
+static std::string addPrefixAndSuffix(const std::string& original_path,
+                                      const std::string& prefix_to_add,
+                                      const std::string& suffix_to_add) {
+  std::string prefixed_path = addPrefixToFileName(original_path, prefix_to_add);
+  return addSuffixBeforeExtension(prefixed_path, suffix_to_add);
+}
+
 PYBIND11_MODULE(_specs, m) {
   auto structs_m = py::module::import("mujoco._structs");
   py::function mjmodel_from_raw_ptr =
@@ -487,11 +510,12 @@ PYBIND11_MODULE(_specs, m) {
             throw pybind11::value_error(mjs_getError(self.ptr));
           }
         }
-        // append the assets adding the suffix to their path
+        // add prefix and suffix to the assets keys
-        std::string suff(s);
+        std::string pre(p);
+        std::string suf(s);
         for (const auto& asset : child.assets) {
           std::string asset_name =
-              addSuffixBeforeExtension(asset.first.cast<std::string>(), suff);
+              addPrefixAndSuffix(asset.first.cast<std::string>(), pre, suf);
           if (self.assets.contains(asset_name) && !self.override_assets) {
             throw pybind11::value_error("Asset " +
                                         asset.first.cast<std::string>() +
@@ -500,22 +524,46 @@ PYBIND11_MODULE(_specs, m) {
           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) {
+        raw::MjsElement* parent_mesh = mjs_firstElement(self.ptr, mjOBJ_MESH);
-          std::string file = mjs_getString(mjs_asTexture(tex)->file);
+        raw::MjsElement* parent_tex = mjs_firstElement(self.ptr, mjOBJ_TEXTURE);
-          if (!file.empty()) {
+        bool child_has_assets = mesh || tex;
-            std::string tex_file = addSuffixBeforeExtension(file, suff);
+        bool parent_has_assets = parent_mesh || parent_tex;
-            mjs_setString(mjs_asTexture(tex)->file, tex_file.c_str());
+        bool child_use_asset_dict = !child.assets.empty();
+        bool parent_use_asset_dict = !self.assets.empty();
+        if (!child_use_asset_dict && child_has_assets &&
+            parent_use_asset_dict) {
+          PyErr_WarnEx(
+              PyExc_Warning,
+              "Attaching a child without asset dict to a parent with an "
+              "asset dict might result in missing assets when attaching again.",
+              1);
+        }
+        if (!parent_use_asset_dict && parent_has_assets &&
+            child_use_asset_dict) {
+          PyErr_WarnEx(
+              PyExc_Warning,
+              "Attaching a child with asset dict to a parent without an "
+              "asset dict might result in missing assets when attaching again.",
+              1);
+        }
+        if (child_use_asset_dict) {
+          while (mesh) {
+            std::string file = mjs_getString(mjs_asMesh(mesh)->file);
+            if (!file.empty()) {
+              std::string mesh_file = addPrefixAndSuffix(file, pre, suf);
+              mjs_setString(mjs_asMesh(mesh)->file, mesh_file.c_str());
+            }
+            mesh = mjs_nextElement(child.ptr, mesh);
+          }
+          while (tex) {
+            std::string file = mjs_getString(mjs_asTexture(tex)->file);
+            if (!file.empty()) {
+              std::string tex_file = addPrefixAndSuffix(file, pre, suf);
+              mjs_setString(mjs_asTexture(tex)->file, tex_file.c_str());
+            }
+            tex = mjs_nextElement(child.ptr, tex);
           }
-          tex = mjs_nextElement(child.ptr, tex);
         }
         child.parent = &self;
         return mjs_asFrame(attached_frame);

python/mujoco/specs_test.py

@@ -1047,7 +1047,7 @@ class SpecsTest(absltest.TestCase):
 
   def test_attach_to_site(self):
     parent = mujoco.MjSpec()
-    parent.assets = {'cube.obj': 'cube_content'}
+    parent.assets = {'path/cube.obj': 'cube_content'}
     site = parent.worldbody.add_site(pos=[1, 2, 3], quat=[0, 0, 0, 1])
     site.name = 'site'
 
@@ -1063,11 +1063,11 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(model1.nbody, 2)
     np.testing.assert_array_equal(model1.body_pos[1], [0, 1, 4])
     np.testing.assert_array_equal(model1.body_quat[1], [0, 0, 0, 1])
-    self.assertEqual(parent.assets['cube.obj'], 'cube_content')
+    self.assertEqual(parent.assets['path/cube.obj'], 'cube_content')
 
     # Attach entire spec to site and compile again.
     child2 = mujoco.MjSpec()
-    child2.assets = {'cube2.obj': 'cube2_content'}
+    child2.assets = {'path/cube2.obj': 'cube2_content'}
     body2 = child2.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child2, site=site, suffix='-child2'))
     self.assertIsNotNone(child2.worldbody)
@@ -1080,14 +1080,14 @@ class SpecsTest(absltest.TestCase):
     np.testing.assert_array_equal(model2.body_pos[2], [2, 3, 2])
     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['path/cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2-child2.obj'], 'cube2_content')
+    self.assertEqual(parent.assets['path/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'}
+    child3.assets = {'path/cube3.obj': 'cube3_content'}
     body3 = child3.worldbody.add_body(name='body')
-    self.assertIsNotNone(parent.attach(child3, site='site', suffix='-child3'))
+    self.assertIsNotNone(parent.attach(child3, site='site', prefix='child3-'))
     self.assertIsNotNone(child3.worldbody)
     self.assertEqual(child3.parent, parent)
     body3.pos = [-2, -2, -2]
@@ -1100,9 +1100,9 @@ class SpecsTest(absltest.TestCase):
     np.testing.assert_array_equal(model3.body_quat[1], [0, 0, 0, 1])
     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['path/cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2-child2.obj'], 'cube2_content')
+    self.assertEqual(parent.assets['path/cube2-child2.obj'], 'cube2_content')
-    self.assertEqual(parent.assets['cube3-child3.obj'], 'cube3_content')
+    self.assertEqual(parent.assets['path/child3-cube3.obj'], 'cube3_content')
 
     # Fail to attach to a site that does not exist.
     child4 = mujoco.MjSpec()
@@ -1127,13 +1127,13 @@ class SpecsTest(absltest.TestCase):
 
   def test_attach_to_frame(self):
     parent = mujoco.MjSpec()
-    parent.assets = {'cube.obj': 'cube_content'}
+    parent.assets = {'path/cube.obj': 'cube_content'}
     frame = parent.worldbody.add_frame(pos=[1, 2, 3], quat=[0, 0, 0, 1])
     frame.name = 'frame'
 
     # Attach body to frame and compile.
     child1 = mujoco.MjSpec()
-    child1.assets = {'cube1.obj': 'cube1_content'}
+    child1.assets = {'path/cube1.obj': 'cube1_content'}
     body1 = child1.worldbody.add_body()
     self.assertIs(body1, frame.attach_body(body1, prefix='_'))
     self.assertIsNotNone(child1.worldbody)
@@ -1143,11 +1143,11 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(model1.nbody, 2)
     np.testing.assert_array_equal(model1.body_pos[1], [0, 1, 4])
     np.testing.assert_array_equal(model1.body_quat[1], [0, 0, 0, 1])
-    self.assertEqual(parent.assets['cube.obj'], 'cube_content')
+    self.assertEqual(parent.assets['path/cube.obj'], 'cube_content')
 
     # Attach entire spec to frame and compile again.
     child2 = mujoco.MjSpec()
-    child2.assets = {'cube2.obj': 'cube2_content'}
+    child2.assets = {'path/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, suffix='-child'))
@@ -1161,15 +1161,15 @@ class SpecsTest(absltest.TestCase):
     np.testing.assert_array_equal(model2.body_pos[2], [3, 4, 3])
     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['path/cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2-child.obj'], 'cube2_content')
+    self.assertEqual(parent.assets['path/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'}
+    child3.assets = {'path/cube2.obj': 'new_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', suffix='-child3'))
+    self.assertIsNotNone(parent.attach(child3, frame='frame', prefix='child3-'))
     self.assertIsNotNone(child3.worldbody)
     self.assertEqual(child3.parent, parent)
     body3.pos = [-2, -2, -2]
@@ -1182,8 +1182,8 @@ class SpecsTest(absltest.TestCase):
     np.testing.assert_array_equal(model3.body_quat[1], [0, 0, 0, 1])
     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['path/cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2-child3.obj'], 'new_cube2_content')
+    self.assertEqual(parent.assets['path/child3-cube2.obj'], 'new_content')
 
     # Fail to attach to a frame that does not exist.
     child4 = mujoco.MjSpec()