From 4903d5321c253a1ad63ea1861ffb5d7a3c54f03d Mon Sep 17 00:00:00 2001 From: Silvia Cruciani Date: Thu, 16 Jan 2025 07:19:03 -0800 Subject: [PATCH] add find_geom property to mjSpec PiperOrigin-RevId: 716225290 Change-Id: Ie6e229b0e802e7df06148f994d598334592e8ad1 --- python/mujoco/specs.cc | 6 ++++++ python/mujoco/specs_test.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 2f37f64f..eb8e0ea9 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -402,6 +402,12 @@ PYBIND11_MODULE(_specs, m) { return mjs_findDefault(self.ptr, classname.c_str()); }, py::return_value_policy::reference_internal); + mjSpec.def( + "find_geom", + [](MjSpec& self, std::string& name) -> raw::MjsGeom* { + return mjs_asGeom(mjs_findElement(self.ptr, mjOBJ_GEOM, name.c_str())); + }, + py::return_value_policy::reference_internal); mjSpec.def("compile", [mjmodel_from_spec_ptr](MjSpec& self) -> py::object { if (self.assets.empty()) { return mjmodel_from_spec_ptr(reinterpret_cast(self.ptr)); diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index d7b86481..993fc2d4 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -691,6 +691,24 @@ class SpecsTest(absltest.TestCase): body4.name = 'body4_new' self.assertEqual(spec.bodies[4].name, 'body4_new') + def test_geom_list(self): + main_xml = """ + + + + + + """ + spec = mujoco.MjSpec.from_string(main_xml) + geom1 = spec.worldbody.add_geom(name='geom1') + geom2 = spec.worldbody.add_geom(name='geom2') + geom3 = spec.find_body('body1').add_geom(name='geom3') + + self.assertEqual(spec.geoms, [geom1, geom2, geom3]) + self.assertEqual(spec.find_geom('geom1'), geom1) + self.assertEqual(spec.find_geom('geom2'), geom2) + self.assertEqual(spec.find_geom('geom3'), geom3) + def test_iterators(self): spec = mujoco.MjSpec() geom1 = spec.worldbody.add_geom()