diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index 78448097..9cdc7034 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -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) diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index e807a8c9..01e9aaf8 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -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( diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 7d94a802..5c9c160e 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -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 = """ + + + + + + + + + + + """ + 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()