Change site.attach to accept a spec instead of a body.

Add site.attach_body for attaching a body to a site.

PiperOrigin-RevId: 703047770
Change-Id: Ibc1f916f8741872135220461cda7cebc4b418913

python/mujoco/specs.cc

@@ -746,7 +746,7 @@ PYBIND11_MODULE(_specs, m) {
       },
       py::return_value_policy::reference_internal);
   mjsSite.def(
-      "attach",
+      "attach_body",
       [](raw::MjsSite& self, raw::MjsBody& body,
          std::optional<std::string>& prefix,
          std::optional<std::string>& suffix) -> raw::MjsBody* {
@@ -762,6 +762,28 @@ PYBIND11_MODULE(_specs, m) {
       py::arg("body"), py::arg("prefix") = py::none(),
       py::arg("suffix") = py::none(),
       py::return_value_policy::reference_internal);
+  mjsSite.def(
+      "attach",
+      [](raw::MjsSite& self, MjSpec& spec,
+         std::optional<std::string>& prefix,
+         std::optional<std::string>& suffix) -> raw::MjsFrame* {
+        auto world = mjs_findBody(spec.ptr, "world");
+        if (!world) {
+          throw pybind11::value_error(
+              mjs_getError(mjs_getSpec(self.element)));
+        }
+        const char* p = prefix.has_value() ? prefix.value().c_str() : "";
+        const char* s = suffix.has_value() ? suffix.value().c_str() : "";
+        auto attached_world = mjs_attachToSite(&self, world, p, s);
+        if (!attached_world) {
+          throw pybind11::value_error(
+              mjs_getError(mjs_getSpec(self.element)));
+        }
+        return mjs_bodyToFrame(&attached_world);
+      },
+      py::arg("body"), py::arg("prefix") = py::none(),
+      py::arg("suffix") = py::none(),
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSCAMERA ===================================
   mjsCamera.def("delete",

python/mujoco/specs_test.py

@@ -875,11 +875,26 @@ class SpecsTest(absltest.TestCase):
   def test_attach_body_to_site(self):
     child = mujoco.MjSpec()
     parent = mujoco.MjSpec()
-    site = parent.worldbody.add_site(pos=[1, 2, 3])
+    site = parent.worldbody.add_site(pos=[1, 2, 3], quat=[0, 0, 0, 1])
     body = child.worldbody.add_body()
-    self.assertIsNotNone(site.attach(body, prefix='_'))
+
-    model = parent.compile()
+    # Attach body to site and compile.
-    np.testing.assert_array_equal(model.body_pos[1], [1, 2, 3])
+    self.assertIsNotNone(site.attach_body(body, prefix='_'))
+    model1 = parent.compile()
+    self.assertIsNotNone(model1)
+    self.assertEqual(model1.nbody, 2)
+    np.testing.assert_array_equal(model1.body_pos[1], [1, 2, 3])
+    np.testing.assert_array_equal(model1.body_quat[1], [0, 0, 0, 1])
+
+    # Attach entire spec to site and compile again.
+    self.assertIsNotNone(site.attach(child, prefix='child-'))
+    model2 = parent.compile()
+    self.assertIsNotNone(model2)
+    self.assertEqual(model2.nbody, 3)
+    np.testing.assert_array_equal(model2.body_pos[1], [1, 2, 3])
+    np.testing.assert_array_equal(model2.body_pos[2], [1, 2, 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])
 
   def test_body_to_frame(self):
     spec = mujoco.MjSpec()
@@ -890,16 +905,27 @@ class SpecsTest(absltest.TestCase):
 
   def test_attach_spec_to_frame(self):
     child = mujoco.MjSpec()
-    child.worldbody.add_camera(name='camera')
     parent = mujoco.MjSpec()
-    frame = parent.worldbody.add_frame(name='frame')
+    frame = parent.worldbody.add_frame(pos=[1, 2, 3], quat=[0, 0, 0, 1])
-    frame.attach(child, prefix='child-')
+    body = child.worldbody.add_body()
-    self.assertLen(child.cameras, 1)
+
-    self.assertLen(parent.bodies, 1)
+    # Attach body to frame and compile.
-    self.assertLen(parent.frames, 2)
+    self.assertIsNotNone(frame.attach_body(body, prefix='_'))
-    self.assertEqual(parent.cameras[0].name, 'child-camera')
+    model1 = parent.compile()
-    self.assertEqual(parent.frames[0].name, 'frame')
+    self.assertIsNotNone(model1)
-    self.assertEqual(parent.frames[1].name, '')
+    self.assertEqual(model1.nbody, 2)
+    np.testing.assert_array_equal(model1.body_pos[1], [1, 2, 3])
+    np.testing.assert_array_equal(model1.body_quat[1], [0, 0, 0, 1])
+
+    # Attach entire spec to frame and compile again.
+    self.assertIsNotNone(frame.attach(child, prefix='child-'))
+    model2 = parent.compile()
+    self.assertIsNotNone(model2)
+    self.assertEqual(model2.nbody, 3)
+    np.testing.assert_array_equal(model2.body_pos[1], [1, 2, 3])
+    np.testing.assert_array_equal(model2.body_pos[2], [1, 2, 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])
 
 
 if __name__ == '__main__':