Add auto-generated tree element lists to Python bindings.
PiperOrigin-RevId: 680973047 Change-Id: I948151260cba330c6b7218cecf13aae2c7ae92c7
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
62ec802ec1
commit
8a825e6050
@@ -252,31 +252,31 @@ def generate() -> None:
|
||||
|
||||
def generate_add() -> None:
|
||||
"""Generate add constructors with optional keyword arguments."""
|
||||
for key, parent, default in [
|
||||
('mjsSite', 'Body', True),
|
||||
('mjsGeom', 'Body', True),
|
||||
('mjsJoint', 'Body', True),
|
||||
('mjsLight', 'Body', True),
|
||||
('mjsCamera', 'Body', True),
|
||||
('mjsBody', 'Body', True),
|
||||
('mjsFrame', 'Body', True),
|
||||
('mjsMaterial', 'Spec', True),
|
||||
('mjsMesh', 'Spec', True),
|
||||
('mjsPair', 'Spec', True),
|
||||
('mjsEquality', 'Spec', True),
|
||||
('mjsTendon', 'Spec', True),
|
||||
('mjsActuator', 'Spec', True),
|
||||
('mjsSkin', 'Spec', False),
|
||||
('mjsTexture', 'Spec', False),
|
||||
('mjsText', 'Spec', False),
|
||||
('mjsTuple', 'Spec', False),
|
||||
('mjsFlex', 'Spec', False),
|
||||
('mjsHField', 'Spec', False),
|
||||
('mjsKey', 'Spec', False),
|
||||
('mjsNumeric', 'Spec', False),
|
||||
('mjsExclude', 'Spec', False),
|
||||
('mjsSensor', 'Spec', False),
|
||||
('mjsPlugin', 'Spec', False),
|
||||
for key, parent, default, listname, objtype in [
|
||||
('mjsSite', 'Body', True, 'sites', 'mjOBJ_SITE'),
|
||||
('mjsGeom', 'Body', True, 'geoms', 'mjOBJ_GEOM'),
|
||||
('mjsJoint', 'Body', True, 'joints', 'mjOBJ_JOINT'),
|
||||
('mjsLight', 'Body', True, 'lights', 'mjOBJ_LIGHT'),
|
||||
('mjsCamera', 'Body', True, 'cameras', 'mjOBJ_CAMERA'),
|
||||
('mjsBody', 'Body', True, 'bodies', 'mjOBJ_BODY'),
|
||||
('mjsFrame', 'Body', True, 'frames', 'mjOBJ_FRAME'),
|
||||
('mjsMaterial', 'Spec', True, 'materials', 'mjOBJ_MATERIAL'),
|
||||
('mjsMesh', 'Spec', True, 'meshes', 'mjOBJ_MESH'),
|
||||
('mjsPair', 'Spec', True, 'pairs', 'mjOBJ_PAIR'),
|
||||
('mjsEquality', 'Spec', True, 'equalities', 'mjOBJ_EQUALITY'),
|
||||
('mjsTendon', 'Spec', True, 'tendons', 'mjOBJ_TENDON'),
|
||||
('mjsActuator', 'Spec', True, 'actuators', 'mjOBJ_ACTUATOR'),
|
||||
('mjsSkin', 'Spec', False, 'skins', 'mjOBJ_SKIN'),
|
||||
('mjsTexture', 'Spec', False, 'textures', 'mjOBJ_TEXTURE'),
|
||||
('mjsText', 'Spec', False, 'texts', 'mjOBJ_TEXT'),
|
||||
('mjsTuple', 'Spec', False, 'tuples', 'mjOBJ_TUPLE'),
|
||||
('mjsFlex', 'Spec', False, 'flexes', 'mjOBJ_FLEX'),
|
||||
('mjsHField', 'Spec', False, 'hfields', 'mjOBJ_HFIELD'),
|
||||
('mjsKey', 'Spec', False, 'keys', 'mjOBJ_KEY'),
|
||||
('mjsNumeric', 'Spec', False, 'numerics', 'mjOBJ_NUMERIC'),
|
||||
('mjsExclude', 'Spec', False, 'excludes', 'mjOBJ_EXCLUDE'),
|
||||
('mjsSensor', 'Spec', False, 'sensors', 'mjOBJ_SENSOR'),
|
||||
('mjsPlugin', 'Spec', False, 'plugins', 'mjOBJ_PLUGIN'),
|
||||
]:
|
||||
|
||||
def _field(f: ast_nodes.StructFieldDecl):
|
||||
@@ -561,6 +561,21 @@ def generate_add() -> None:
|
||||
py::return_value_policy::reference_internal);
|
||||
"""
|
||||
|
||||
code += f"""\n
|
||||
mjSpec.def_property_readonly(
|
||||
"{listname}",
|
||||
[](MjSpec& self) -> py::list {{
|
||||
py::list list;
|
||||
raw::MjsElement* el = mjs_firstElement(self.ptr, {objtype});
|
||||
while (el) {{
|
||||
list.append(mjs_as{key[3:]}(el));
|
||||
el = mjs_nextElement(self.ptr, el);
|
||||
}}
|
||||
return list;
|
||||
}},
|
||||
py::return_value_policy::reference_internal);
|
||||
"""
|
||||
|
||||
print(code)
|
||||
|
||||
|
||||
|
||||
@@ -333,210 +333,6 @@ PYBIND11_MODULE(_specs, m) {
|
||||
mjSpec.def("detach_body", [](MjSpec& self, raw::MjsBody& body) {
|
||||
mjs_detachBody(self.ptr, &body);
|
||||
});
|
||||
mjSpec.def_property_readonly(
|
||||
"plugins",
|
||||
[](MjSpec& self) -> py::list {
|
||||
py::list list;
|
||||
raw::MjsElement* el = mjs_firstElement(self.ptr, mjOBJ_PLUGIN);
|
||||
while (el) {
|
||||
list.append(mjs_asPlugin(el));
|
||||
el = mjs_nextElement(self.ptr, el);
|
||||
}
|
||||
return list;
|
||||
},
|
||||
py::return_value_policy::reference_internal);
|
||||
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_property_readonly(
|
||||
|
||||
@@ -574,6 +574,26 @@ class SpecsTest(absltest.TestCase):
|
||||
self.assertEqual(spec.sensors[1].name, 'sensor2')
|
||||
self.assertEqual(spec.sensors[2].name, 'sensor3')
|
||||
|
||||
def test_body_list(self):
|
||||
main_xml = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="body1">
|
||||
<body name="body3">
|
||||
<body name="body4"/>
|
||||
</body>
|
||||
</body>
|
||||
<body name="body2"/>
|
||||
</worldbody>
|
||||
</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')
|
||||
|
||||
def test_iterators(self):
|
||||
spec = mujoco.MjSpec()
|
||||
geom1 = spec.worldbody.add_geom()
|
||||
|
||||
Reference in New Issue
Block a user