add find_geom property to mjSpec

PiperOrigin-RevId: 716225290
Change-Id: Ie6e229b0e802e7df06148f994d598334592e8ad1

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<uintptr_t>(self.ptr));

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 = """
+    <mujoco>
+      <worldbody>
+        <body name="body1"/>
+      </worldbody>
+    </mujoco>
+    """
+    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()