Add non-tree element lists to python bindings.

Fixes #1752.

PiperOrigin-RevId: 649405039
Change-Id: Ie0b575a3ec912b773141694397ac2f394bff1227
This commit is contained in:
Alessio Quaglino
2024-07-04 07:23:05 -07:00
committed by Copybara-Service
parent d831153401
commit a0945e4b5a
10 changed files with 984 additions and 0 deletions
+192
View File
@@ -297,6 +297,198 @@ PYBIND11_MODULE(_specs, m) {
mjSpec.def("detach_body", [](MjSpec& self, raw::MjsBody& body) {
mjs_detachBody(self.ptr, &body);
});
mjSpec.def_property_readonly(
"actuators",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_ACTUATOR);
while (el) {
list.append(mjs_asActuator(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"sensors",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_SENSOR);
while (el) {
list.append(mjs_asSensor(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"flexes",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_FLEX);
while (el) {
list.append(mjs_asFlex(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"pairs",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_PAIR);
while (el) {
list.append(mjs_asPair(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"equality",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_EQUALITY);
while (el) {
list.append(mjs_asEquality(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"excludes",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_EXCLUDE);
while (el) {
list.append(mjs_asExclude(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"tendons",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_TENDON);
while (el) {
list.append(mjs_asTendon(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"numeric",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_NUMERIC);
while (el) {
list.append(mjs_asNumeric(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"text",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_TEXT);
while (el) {
list.append(mjs_asText(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"tuple",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_TUPLE);
while (el) {
list.append(mjs_asTuple(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"key",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_KEY);
while (el) {
list.append(mjs_asKey(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"mesh",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_MESH);
while (el) {
list.append(mjs_asMesh(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"hfield",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_HFIELD);
while (el) {
list.append(mjs_asHField(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"skin",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_SKIN);
while (el) {
list.append(mjs_asSkin(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"texture",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_TEXTURE);
while (el) {
list.append(mjs_asTexture(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
mjSpec.def_property_readonly(
"material",
[](MjSpec& self) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_MATERIAL);
while (el) {
list.append(mjs_asMaterial(el));
el = mjs_nextElement(self.ptr, el);
}
return list;
},
py::return_value_policy::reference_internal);
// ============================= MJSBODY ====================================
mjsBody.def("id", [](raw::MjsBody& self) -> int {
+13
View File
@@ -217,6 +217,19 @@ class SpecsTest(absltest.TestCase):
</mujoco>
"""))
def test_element_list(self):
spec = mujoco.MjSpec()
sensor1 = spec.add_sensor()
sensor2 = spec.add_sensor()
sensor3 = spec.add_sensor()
sensor1.name = 'sensor1'
sensor2.name = 'sensor2'
sensor3.name = 'sensor3'
self.assertLen(spec.sensors, 3)
self.assertEqual(spec.sensors[0].name, 'sensor1')
self.assertEqual(spec.sensors[1].name, 'sensor2')
self.assertEqual(spec.sensors[2].name, 'sensor3')
if __name__ == '__main__':
absltest.main()