Add sibling traversal for tree elements.

PiperOrigin-RevId: 691057957
Change-Id: I3f30cf175b5ad0842f3fbbf7fdb42ba7a2212682

python/mujoco/specs.cc

@@ -155,9 +155,9 @@ void DefineArray(py::module& m, const std::string& typestr) {
       }, py::keep_alive<0, 1>(), py::return_value_policy::reference_internal);
 };
 
-py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype) {
+py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) {
   py::list list;
-  raw::MjsElement* el = mjs_firstChild(&body, objtype, true);
+  raw::MjsElement* el = mjs_firstChild(&body, objtype, recursive);
   std::string error = mjs_getError(mjs_getSpec(body.element));
   if (!el && !error.empty()) {
     throw pybind11::value_error(error);
@@ -192,7 +192,7 @@ py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype) {
             "light, camera.");
         break;
     }
-    el = mjs_nextChild(&body, el, true);
+    el = mjs_nextChild(&body, el, recursive);
   }
   return list;  // list of pointers, so they can be copied
 }
@@ -460,7 +460,7 @@ PYBIND11_MODULE(_specs, m) {
   mjsBody.def(
       "find_all",
       [](raw::MjsBody& self, mjtObj objtype) -> py::list {
-        return FindAllImpl(self, objtype);
+        return FindAllImpl(self, objtype, true);
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
@@ -484,7 +484,7 @@ PYBIND11_MODULE(_specs, m) {
               "body.find_all supports the types: body, frame, geom, site, "
               "light, camera.");
         }
-        return FindAllImpl(self, objtype);
+        return FindAllImpl(self, objtype, true);
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
@@ -505,6 +505,12 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_asBody(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def_property_readonly(
+      "bodies",
+      [](raw::MjsBody& self) -> py::list {
+        return FindAllImpl(self, mjOBJ_BODY, false);
+      },
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_camera",
       [](raw::MjsBody& self) -> raw::MjsCamera* {
@@ -517,6 +523,12 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_asCamera(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def_property_readonly(
+      "cameras",
+      [](raw::MjsBody& self) -> py::list {
+        return FindAllImpl(self, mjOBJ_CAMERA, false);
+      },
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_light",
       [](raw::MjsBody& self) -> raw::MjsLight* {
@@ -529,6 +541,12 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_asLight(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def_property_readonly(
+      "lights",
+      [](raw::MjsBody& self) -> py::list {
+        return FindAllImpl(self, mjOBJ_LIGHT, false);
+      },
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_joint",
       [](raw::MjsBody& self) -> raw::MjsJoint* {
@@ -541,6 +559,12 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_asJoint(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def_property_readonly(
+      "joints",
+      [](raw::MjsBody& self) -> py::list {
+        return FindAllImpl(self, mjOBJ_JOINT, false);
+      },
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_geom",
       [](raw::MjsBody& self) -> raw::MjsGeom* {
@@ -553,6 +577,12 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_asGeom(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def_property_readonly(
+      "geoms",
+      [](raw::MjsBody& self) -> py::list {
+        return FindAllImpl(self, mjOBJ_GEOM, false);
+      },
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_site",
       [](raw::MjsBody& self) -> raw::MjsSite* {
@@ -565,6 +595,12 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_asSite(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def_property_readonly(
+      "sites",
+      [](raw::MjsBody& self) -> py::list {
+        return FindAllImpl(self, mjOBJ_SITE, false);
+      },
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_frame",
       [](raw::MjsBody& self) -> raw::MjsFrame* {
@@ -577,6 +613,12 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_asFrame(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def_property_readonly(
+      "frames",
+      [](raw::MjsBody& self) -> py::list {
+        return FindAllImpl(self, mjOBJ_FRAME, false);
+      },
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "spec",
       [](raw::MjsBody& self) -> raw::MjSpec* {

python/mujoco/specs_test.py

@@ -602,8 +602,12 @@ class SpecsTest(absltest.TestCase):
       <worldbody>
         <body name="body1">
           <body name="body3">
+            <site name="site1"/>
+            <site name="site2"/>
+            <site name="site3"/>
+            <site name="site4"/>
             <body name="body4">
-              <site name="site"/>
+              <site name="site5"/>
             </body>
           </body>
         </body>
@@ -614,6 +618,9 @@ class SpecsTest(absltest.TestCase):
     spec = mujoco.MjSpec.from_string(main_xml)
     bodytype = mujoco.mjtObj.mjOBJ_BODY
     self.assertLen(spec.bodies, 5)
+    self.assertLen(spec.sites, 5)
+    self.assertLen(spec.worldbody.find_all('body'), 4)
+    self.assertLen(spec.worldbody.find_all('site'), 5)
     self.assertEqual(spec.bodies[1].name, 'body1')
     self.assertEqual(spec.bodies[2].name, 'body2')
     self.assertEqual(spec.bodies[3].name, 'body3')
@@ -630,7 +637,19 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(spec.bodies[3].find_all('body')[0].name, 'body4')
     self.assertEmpty(spec.bodies[2].find_all('body'))
     self.assertEmpty(spec.bodies[4].find_all('body'))
-    self.assertEqual(spec.worldbody.find_all('site')[0].name, 'site')
+    self.assertEqual(spec.worldbody.find_all('site')[0].name, 'site1')
+    self.assertEqual(spec.worldbody.find_all('site')[1].name, 'site2')
+    self.assertEqual(spec.worldbody.find_all('site')[2].name, 'site3')
+    self.assertEqual(spec.worldbody.find_all('site')[3].name, 'site4')
+    self.assertEqual(spec.worldbody.find_all('site')[4].name, 'site5')
+    self.assertEmpty(spec.bodies[2].sites)
+    self.assertLen(spec.bodies[3].sites, 4)
+    self.assertLen(spec.bodies[4].sites, 1)
+    self.assertEqual(spec.bodies[3].sites[0].name, 'site1')
+    self.assertEqual(spec.bodies[3].sites[1].name, 'site2')
+    self.assertEqual(spec.bodies[3].sites[2].name, 'site3')
+    self.assertEqual(spec.bodies[3].sites[3].name, 'site4')
+    self.assertEqual(spec.bodies[4].sites[0].name, 'site5')
     with self.assertRaises(ValueError) as cm:
       spec.worldbody.find_all('actuator')
     self.assertEqual(