diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index fe7e22fb..b2beaa3f 100644 --- a/python/mujoco/specs.cc +++ b/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* { diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index fc325a4a..7bca48f0 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -602,8 +602,12 @@ class SpecsTest(absltest.TestCase): + + + + - + @@ -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(