Add body.find_all() to Python bindings.

This enables to generate a list of all body children of a given type.

Also, add recursive option to mjs_firstChild and mjs_nextChild. Fixes #2112.

PiperOrigin-RevId: 681039348
Change-Id: Iebbb26fd5ec314150206edfdddd88633baa0bec6

doc/APIreference/functions.rst

@@ -4205,7 +4205,7 @@ mjs_firstChild
 
 .. mujoco-include:: mjs_firstChild
 
-Return body's first child of given type.
+Return body's first child of given type. If recurse is nonzero, also search the body's subtree.
 
 .. _mjs_nextChild:
 
@@ -4215,6 +4215,7 @@ mjs_nextChild
 .. mujoco-include:: mjs_nextChild
 
 Return body's next child of the same type; return NULL if child is last.
+If recurse is nonzero, also search the body's subtree.
 
 .. _mjs_firstElement:
 

doc/includes/references.h

@@ -3598,8 +3598,8 @@ mjsDefault* mjs_getDefault(mjsElement* element);
 mjsDefault* mjs_findDefault(mjSpec* s, const char* classname);
 mjsDefault* mjs_getSpecDefault(mjSpec* s);
 int mjs_getId(mjsElement* element);
-mjsElement* mjs_firstChild(mjsBody* body, mjtObj type);
+mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse);
-mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child);
+mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse);
 mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);
 mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element);
 void mjs_setBuffer(mjByteVec* dest, const void* array, int size);

include/mujoco/mujoco.h

@@ -1550,11 +1550,12 @@ MJAPI mjsDefault* mjs_getSpecDefault(mjSpec* s);
 // Get element id.
 MJAPI int mjs_getId(mjsElement* element);
 
-// Return body's first child of given type.
+// Return body's first child of given type. If recurse is nonzero, also search the body's subtree.
-MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type);
+MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse);
 
 // Return body's next child of the same type; return NULL if child is last.
-MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child);
+// If recurse is nonzero, also search the body's subtree.
+MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse);
 
 // Return spec's first element of selected type.
 MJAPI mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);

introspect/functions.py

@@ -9859,8 +9859,12 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  name='type',
                  type=ValueType(name='mjtObj'),
              ),
+             FunctionParameterDecl(
+                 name='recurse',
+                 type=ValueType(name='int'),
+             ),
          ),
-         doc="Return body's first child of given type.",
+         doc="Return body's first child of given type. If recurse is nonzero, also search the body's subtree.",  # pylint: disable=line-too-long
      )),
     ('mjs_nextChild',
      FunctionDecl(
@@ -9881,8 +9885,12 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                      inner_type=ValueType(name='mjsElement'),
                  ),
              ),
+             FunctionParameterDecl(
+                 name='recurse',
+                 type=ValueType(name='int'),
+             ),
          ),
-         doc="Return body's next child of the same type; return NULL if child is last.",  # pylint: disable=line-too-long
+         doc="Return body's next child of the same type; return NULL if child is last. If recurse is nonzero, also search the body's subtree.",  # pylint: disable=line-too-long
      )),
     ('mjs_firstElement',
      FunctionDecl(

python/mujoco/specs.cc

@@ -385,6 +385,50 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_getDefault(self.element);
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def(
+      "find_all",
+      [](raw::MjsBody& self, mjtObj objtype) -> py::list {
+        py::list list;
+        raw::MjsElement* el = mjs_firstChild(&self, objtype, true);
+        std::string error = mjs_getError(mjs_getSpec(&self));
+        if (!el && !error.empty()) {
+          throw pybind11::value_error(error);
+        }
+        while (el) {
+          switch (objtype) {
+            case mjOBJ_BODY:
+              list.append(mjs_asBody(el));
+              break;
+            case mjOBJ_CAMERA:
+              list.append(mjs_asCamera(el));
+              break;
+            case mjOBJ_FRAME:
+              list.append(mjs_asFrame(el));
+              break;
+            case mjOBJ_GEOM:
+              list.append(mjs_asGeom(el));
+              break;
+            case mjOBJ_JOINT:
+              list.append(mjs_asJoint(el));
+              break;
+            case mjOBJ_LIGHT:
+              list.append(mjs_asLight(el));
+              break;
+            case mjOBJ_SITE:
+              list.append(mjs_asSite(el));
+              break;
+            default:
+              // this should never happen
+              throw pybind11::value_error(
+                  "body.find_all supports the types: body, frame, geom, site, "
+                  "light, camera.");
+              break;
+          }
+          el = mjs_nextChild(&self, el, true);
+        }
+        return list;
+      },
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "find_child",
       [](raw::MjsBody& self, std::string& name) -> raw::MjsBody* {
@@ -394,85 +438,85 @@ PYBIND11_MODULE(_specs, m) {
   mjsBody.def(
       "first_body",
       [](raw::MjsBody& self) -> raw::MjsBody* {
-        return mjs_asBody(mjs_firstChild(&self, mjOBJ_BODY));
+        return mjs_asBody(mjs_firstChild(&self, mjOBJ_BODY, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "next_body",
       [](raw::MjsBody& self, raw::MjsBody& child) -> raw::MjsBody* {
-        return mjs_asBody(mjs_nextChild(&self, child.element));
+        return mjs_asBody(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_camera",
       [](raw::MjsBody& self) -> raw::MjsCamera* {
-        return mjs_asCamera(mjs_firstChild(&self, mjOBJ_CAMERA));
+        return mjs_asCamera(mjs_firstChild(&self, mjOBJ_CAMERA, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "next_camera",
       [](raw::MjsBody& self, raw::MjsCamera& child) -> raw::MjsCamera* {
-        return mjs_asCamera(mjs_nextChild(&self, child.element));
+        return mjs_asCamera(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_light",
       [](raw::MjsBody& self) -> raw::MjsLight* {
-        return mjs_asLight(mjs_firstChild(&self, mjOBJ_LIGHT));
+        return mjs_asLight(mjs_firstChild(&self, mjOBJ_LIGHT, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "next_light",
       [](raw::MjsBody& self, raw::MjsLight& child) -> raw::MjsLight* {
-        return mjs_asLight(mjs_nextChild(&self, child.element));
+        return mjs_asLight(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_joint",
       [](raw::MjsBody& self) -> raw::MjsJoint* {
-        return mjs_asJoint(mjs_firstChild(&self, mjOBJ_JOINT));
+        return mjs_asJoint(mjs_firstChild(&self, mjOBJ_JOINT, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "next_joint",
       [](raw::MjsBody& self, raw::MjsJoint& child) -> raw::MjsJoint* {
-        return mjs_asJoint(mjs_nextChild(&self, child.element));
+        return mjs_asJoint(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_geom",
       [](raw::MjsBody& self) -> raw::MjsGeom* {
-        return mjs_asGeom(mjs_firstChild(&self, mjOBJ_GEOM));
+        return mjs_asGeom(mjs_firstChild(&self, mjOBJ_GEOM, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "next_geom",
       [](raw::MjsBody& self, raw::MjsGeom& child) -> raw::MjsGeom* {
-        return mjs_asGeom(mjs_nextChild(&self, child.element));
+        return mjs_asGeom(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_site",
       [](raw::MjsBody& self) -> raw::MjsSite* {
-        return mjs_asSite(mjs_firstChild(&self, mjOBJ_SITE));
+        return mjs_asSite(mjs_firstChild(&self, mjOBJ_SITE, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "next_site",
       [](raw::MjsBody& self, raw::MjsSite& child) -> raw::MjsSite* {
-        return mjs_asSite(mjs_nextChild(&self, child.element));
+        return mjs_asSite(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "first_frame",
       [](raw::MjsBody& self) -> raw::MjsFrame* {
-        return mjs_asFrame(mjs_firstChild(&self, mjOBJ_FRAME));
+        return mjs_asFrame(mjs_firstChild(&self, mjOBJ_FRAME, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "next_frame",
       [](raw::MjsBody& self, raw::MjsFrame& child) -> raw::MjsFrame* {
-        return mjs_asFrame(mjs_nextChild(&self, child.element));
+        return mjs_asFrame(mjs_nextChild(&self, child.element, false));
       },
       py::return_value_policy::reference_internal);
   mjsBody.def(

python/mujoco/specs_test.py

@@ -580,7 +580,9 @@ class SpecsTest(absltest.TestCase):
       <worldbody>
         <body name="body1">
           <body name="body3">
-            <body name="body4"/>
+            <body name="body4">
+              <site name="site"/>
+            </body>
           </body>
         </body>
         <body name="body2"/>
@@ -588,11 +590,33 @@ class SpecsTest(absltest.TestCase):
     </mujoco>
     """
     spec = mujoco.MjSpec.from_string(main_xml)
-    self.assertLen(spec.bodies, 4)
+    bodytype = mujoco.mjtObj.mjOBJ_BODY
-    self.assertEqual(spec.bodies[0].name, 'body1')
+    sitetype = mujoco.mjtObj.mjOBJ_SITE
-    self.assertEqual(spec.bodies[1].name, 'body2')
+    self.assertLen(spec.bodies, 5)
-    self.assertEqual(spec.bodies[2].name, 'body3')
+    self.assertEqual(spec.bodies[1].name, 'body1')
-    self.assertEqual(spec.bodies[3].name, 'body4')
+    self.assertEqual(spec.bodies[2].name, 'body2')
+    self.assertEqual(spec.bodies[3].name, 'body3')
+    self.assertEqual(spec.bodies[4].name, 'body4')
+    self.assertLen(spec.worldbody.find_all(bodytype), 4)
+    self.assertLen(spec.bodies[1].find_all(bodytype), 2)
+    self.assertLen(spec.bodies[3].find_all(bodytype), 1)
+    self.assertEqual(spec.worldbody.find_all(bodytype)[0].name, 'body1')
+    self.assertEqual(spec.worldbody.find_all(bodytype)[1].name, 'body2')
+    self.assertEqual(spec.worldbody.find_all(bodytype)[2].name, 'body3')
+    self.assertEqual(spec.worldbody.find_all(bodytype)[3].name, 'body4')
+    self.assertEqual(spec.bodies[1].find_all(bodytype)[0].name, 'body3')
+    self.assertEqual(spec.bodies[1].find_all(bodytype)[1].name, 'body4')
+    self.assertEqual(spec.bodies[3].find_all(bodytype)[0].name, 'body4')
+    self.assertEmpty(spec.bodies[2].find_all(bodytype))
+    self.assertEmpty(spec.bodies[4].find_all(bodytype))
+    self.assertEqual(spec.worldbody.find_all(sitetype)[0].name, 'site')
+    with self.assertRaises(ValueError) as cm:
+      spec.worldbody.find_all(mujoco.mjtObj.mjOBJ_ACTUATOR)
+    self.assertEqual(
+        str(cm.exception),
+        'Error: Body.NextChild supports the types: body, frame, geom, site,'
+        ' light, camera\nElement name \'world\', id 0',
+    )
 
   def test_iterators(self):
     spec = mujoco.MjSpec()

src/user/user_api.cc

@@ -655,17 +655,27 @@ void mjs_setDefault(mjsElement* element, mjsDefault* defspec) {
 
 
 // return first child of selected type
-mjsElement* mjs_firstChild(mjsBody* body, mjtObj type) {
+mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse) {
   mjCBody* bodyC = static_cast<mjCBody*>(body->element);
-  return bodyC->NextChild(NULL, type);
+  try {
+    return bodyC->NextChild(NULL, type, recurse);
+  } catch (mjCError& e) {
+    bodyC->model->SetError(e);
+    return nullptr;
+  }
 }
 
 
 
 // return body's next child; return NULL if child is last
-mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child) {
+mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse) {
   mjCBody* bodyC = static_cast<mjCBody*>(body->element);
-  return bodyC->NextChild(child);
+  try {
+    return bodyC->NextChild(child, child->elemtype, recurse);
+  } catch(mjCError& e) {
+    bodyC->model->SetError(e);
+    return nullptr;
+  }
 }
 
 

src/user/user_api.h

@@ -218,11 +218,12 @@ MJAPI int mjs_getId(mjsElement* element);
 
 //---------------------------------- Tree traversal ------------------------------------------------
 
-// Return body's first child of given type.
+// Return body's first child of given type. If recurse is nonzero, also search the body's subtree.
-MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type);
+MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse);
 
 // Return body's next child of the same type; return NULL if child is last.
-MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child);
+// If recurse is nonzero, also search the body's subtree.
+MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse);
 
 // Return spec's first element of selected type.
 MJAPI mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);

src/user/user_model.cc

@@ -921,6 +921,13 @@ mjsElement* mjCModel::NextObject(mjsElement* object, mjtObj type) {
 
   switch (type) {
     case mjOBJ_BODY:
+      if (!object) {
+        return bodies_[0]->spec.element;
+      } else if (object == bodies_[0]->spec.element) {
+        return bodies_[0]->NextChild(NULL, type, /*recursive=*/true);
+      } else {
+        return bodies_[0]->NextChild(object, type, /*recursive=*/true);
+      }
     case mjOBJ_SITE:
     case mjOBJ_GEOM:
     case mjOBJ_JOINT:

src/user/user_objects.cc

@@ -1338,6 +1338,9 @@ mjsElement* mjCBody::NextChild(mjsElement* child, mjtObj type, bool recursive) {
       candidate = GetNext(frames, child, recursive);
       break;
     default:
+      throw mjCError(this,
+                     "Body.NextChild supports the types: body, frame, geom, "
+                     "site, light, camera");
       break;
   }
 

test/user/user_api_test.cc

@@ -88,20 +88,25 @@ TEST_F(MujocoTest, TreeTraversal) {
   mjs_setString(geom3->name, "geom3");
   mjs_setString(site4->name, "site4");
 
+  bool recursive = false;
+  mjsElement* b_el0 = mjs_firstElement(spec, mjOBJ_BODY);
+  mjsElement* b_el1 = mjs_nextElement(spec, b_el0);
   mjsElement* a_el0 = mjs_firstElement(spec, mjOBJ_ACTUATOR);
   mjsElement* l_el0 = mjs_firstElement(spec, mjOBJ_LIGHT);
-  mjsElement* c_el0 = mjs_firstChild(body, mjOBJ_CAMERA);
+  mjsElement* c_el0 = mjs_firstChild(body, mjOBJ_CAMERA, recursive);
-  mjsElement* t_el0 = mjs_firstChild(body, mjOBJ_TENDON);
+  mjsElement* t_el0 = mjs_firstChild(body, mjOBJ_TENDON, recursive);
-  mjsElement* s_el1 = mjs_firstChild(body, mjOBJ_SITE);
+  mjsElement* s_el1 = mjs_firstChild(body, mjOBJ_SITE, recursive);
-  mjsElement* s_el2 = mjs_nextChild(body, s_el1);
+  mjsElement* s_el2 = mjs_nextChild(body, s_el1, recursive);
-  mjsElement* s_el3 = mjs_nextChild(body, s_el2);
+  mjsElement* s_el3 = mjs_nextChild(body, s_el2, recursive);
-  mjsElement* s_el0 = mjs_nextChild(body, s_el3);
+  mjsElement* s_el0 = mjs_nextChild(body, s_el3, recursive);
-  mjsElement* s_el4 = mjs_firstChild(body1, mjOBJ_SITE);
+  mjsElement* s_el4 = mjs_firstChild(body1, mjOBJ_SITE, recursive);
-  mjsElement* g_el1 = mjs_firstChild(body, mjOBJ_GEOM);
+  mjsElement* g_el1 = mjs_firstChild(body, mjOBJ_GEOM, recursive);
-  mjsElement* g_el2 = mjs_nextChild(body, g_el1);
+  mjsElement* g_el2 = mjs_nextChild(body, g_el1, recursive);
-  mjsElement* g_el3 = mjs_nextChild(body, g_el2);
+  mjsElement* g_el3 = mjs_nextChild(body, g_el2, recursive);
-  mjsElement* g_el0 = mjs_nextChild(body, g_el3);
+  mjsElement* g_el0 = mjs_nextChild(body, g_el3, recursive);
 
+  EXPECT_EQ(b_el0, world->element);
+  EXPECT_EQ(b_el1, body->element);
   EXPECT_EQ(a_el0, nullptr);
   EXPECT_EQ(l_el0, nullptr);
   EXPECT_EQ(c_el0, nullptr);