chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add sibling traversal for tree elements.

Browse Source 
PiperOrigin-RevId: 691057957
Change-Id: I3f30cf175b5ad0842f3fbbf7fdb42ba7a2212682
This commit is contained in:
Alessio Quaglino
2024-10-29 09:54:07 -07:00
committed by Copybara-Service
parent b941e994d9
commit e835f139c8
2 changed files with 68 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs.cc
+47 -5
View File
@@ -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* {