Add parent to MjSpec wrapper in Python bindings.

This keeps track of specs attached with the full attach function `attach`. It does not track fine-grained attach via `attach_body` and `attach_frame`.

PiperOrigin-RevId: 722597081
Change-Id: I56188f1718592128b7033538acc9a40553cbbed2

python/mujoco/specs.cc

@@ -83,6 +83,7 @@ struct MjSpec {
     for (const auto [key, value] : other.assets) {
       assets[key] = value;
     }
+    parent = other.parent;
   }
   MjSpec& operator=(const MjSpec& other) {
     override_assets = other.override_assets;
@@ -90,6 +91,7 @@ struct MjSpec {
     for (const auto [key, value] : other.assets) {
       assets[key] = value;
     }
+    parent = other.parent;
     return *this;
   }
 
@@ -101,6 +103,8 @@ struct MjSpec {
       assets[key] = value;
     }
     other.assets.clear();
+    parent = other.parent;
+    other.parent = nullptr;
   }
   MjSpec& operator=(MjSpec&& other) {
     override_assets = other.override_assets;
@@ -110,6 +114,8 @@ struct MjSpec {
       assets[key] = value;
     }
     other.assets.clear();
+    parent = other.parent;
+    other.parent = nullptr;
     return *this;
   }
 
@@ -119,6 +125,7 @@ struct MjSpec {
   raw::MjSpec* ptr;
   py::dict assets;
   bool override_assets = true;
+  MjSpec* parent = nullptr;
 };
 
 template <typename LoadFunc>
@@ -272,6 +279,8 @@ PYBIND11_MODULE(_specs, m) {
 
   // ============================= MJSPEC =====================================
   mjSpec.def(py::init<>());
+  mjSpec.def_property_readonly(
+      "parent", [](MjSpec& self) -> MjSpec* { return self.parent; });
   mjSpec.def_static(
       "from_file",
       [](std::string& filename,
@@ -590,6 +599,7 @@ PYBIND11_MODULE(_specs, m) {
           }
           self.assets[asset.first] = asset.second;
         }
+        child.parent = &self;
         return mjs_bodyToFrame(&attached_world);
       },
       py::arg("child"), py::arg("prefix") = py::none(),

python/mujoco/specs_test.py

@@ -981,6 +981,7 @@ class SpecsTest(absltest.TestCase):
     body2 = child2.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child2, site=site, prefix='child2-'))
     self.assertIsNone(child2.worldbody)
+    self.assertEqual(child2.parent, parent)
     body2.pos = [-1, -1, -1]
     model2 = parent.compile()
     self.assertIsNotNone(model2)
@@ -998,6 +999,7 @@ class SpecsTest(absltest.TestCase):
     body3 = child3.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child3, site='site', prefix='child3-'))
     self.assertIsNone(child3.worldbody)
+    self.assertEqual(child3.parent, parent)
     body3.pos = [-2, -2, -2]
     model3 = parent.compile()
     self.assertIsNotNone(model3)
@@ -1050,6 +1052,7 @@ class SpecsTest(absltest.TestCase):
     body2 = child2.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child2, frame=frame, prefix='child-'))
     self.assertIsNone(child2.worldbody)
+    self.assertEqual(child2.parent, parent)
     body2.pos = [-1, -1, -1]
     model2 = parent.compile()
     self.assertIsNotNone(model2)
@@ -1067,6 +1070,7 @@ class SpecsTest(absltest.TestCase):
     body3 = child3.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child3, frame='frame', prefix='child3-'))
     self.assertIsNone(child3.worldbody)
+    self.assertEqual(child3.parent, parent)
     body3.pos = [-2, -2, -2]
     model3 = parent.compile()
     self.assertIsNotNone(model3)