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
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
8a825e6050
commit
89aea5e217
+58
-14
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user