diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index ed317aea..fbec781a 100644 --- a/doc/APIreference/functions.rst +++ b/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: diff --git a/doc/includes/references.h b/doc/includes/references.h index 5f5ecafd..19ea8a36 100644 --- a/doc/includes/references.h +++ b/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_nextChild(mjsBody* body, mjsElement* child); +mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse); +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); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index e9a41e0d..0df9b7e9 100644 --- a/include/mujoco/mujoco.h +++ b/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. -MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj 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, 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); diff --git a/introspect/functions.py b/introspect/functions.py index 08f3d44c..25f21010 100644 --- a/introspect/functions.py +++ b/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( diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 01e9aaf8..8a1e7ecc 100644 --- a/python/mujoco/specs.cc +++ b/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( diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 5c9c160e..68e0f0cb 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -580,7 +580,9 @@ class SpecsTest(absltest.TestCase): - + + + @@ -588,11 +590,33 @@ class SpecsTest(absltest.TestCase): """ spec = mujoco.MjSpec.from_string(main_xml) - self.assertLen(spec.bodies, 4) - self.assertEqual(spec.bodies[0].name, 'body1') - self.assertEqual(spec.bodies[1].name, 'body2') - self.assertEqual(spec.bodies[2].name, 'body3') - self.assertEqual(spec.bodies[3].name, 'body4') + bodytype = mujoco.mjtObj.mjOBJ_BODY + sitetype = mujoco.mjtObj.mjOBJ_SITE + self.assertLen(spec.bodies, 5) + self.assertEqual(spec.bodies[1].name, 'body1') + 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() diff --git a/src/user/user_api.cc b/src/user/user_api.cc index a401f131..55472226 100644 --- a/src/user/user_api.cc +++ b/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(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(body->element); - return bodyC->NextChild(child); + try { + return bodyC->NextChild(child, child->elemtype, recurse); + } catch(mjCError& e) { + bodyC->model->SetError(e); + return nullptr; + } } diff --git a/src/user/user_api.h b/src/user/user_api.h index 43cff0c5..8c55c0ff 100644 --- a/src/user/user_api.h +++ b/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. -MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj 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, 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); diff --git a/src/user/user_model.cc b/src/user/user_model.cc index a69dfbc4..1a6bfbd7 100644 --- a/src/user/user_model.cc +++ b/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: diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 35bd3645..833a8106 100644 --- a/src/user/user_objects.cc +++ b/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; } diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 6e7fd45f..a4a7d425 100644 --- a/test/user/user_api_test.cc +++ b/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* t_el0 = mjs_firstChild(body, mjOBJ_TENDON); - mjsElement* s_el1 = mjs_firstChild(body, mjOBJ_SITE); - mjsElement* s_el2 = mjs_nextChild(body, s_el1); - mjsElement* s_el3 = mjs_nextChild(body, s_el2); - mjsElement* s_el0 = mjs_nextChild(body, s_el3); - mjsElement* s_el4 = mjs_firstChild(body1, mjOBJ_SITE); - mjsElement* g_el1 = mjs_firstChild(body, mjOBJ_GEOM); - mjsElement* g_el2 = mjs_nextChild(body, g_el1); - mjsElement* g_el3 = mjs_nextChild(body, g_el2); - mjsElement* g_el0 = mjs_nextChild(body, g_el3); + mjsElement* c_el0 = mjs_firstChild(body, mjOBJ_CAMERA, recursive); + mjsElement* t_el0 = mjs_firstChild(body, mjOBJ_TENDON, recursive); + mjsElement* s_el1 = mjs_firstChild(body, mjOBJ_SITE, recursive); + mjsElement* s_el2 = mjs_nextChild(body, s_el1, recursive); + mjsElement* s_el3 = mjs_nextChild(body, s_el2, recursive); + mjsElement* s_el0 = mjs_nextChild(body, s_el3, recursive); + mjsElement* s_el4 = mjs_firstChild(body1, mjOBJ_SITE, recursive); + mjsElement* g_el1 = mjs_firstChild(body, mjOBJ_GEOM, recursive); + mjsElement* g_el2 = mjs_nextChild(body, g_el1, recursive); + mjsElement* g_el3 = mjs_nextChild(body, g_el2, recursive); + 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);