From a3b9f986c04c5e353e7546b4e0b0b1ecd48e7395 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Wed, 4 Dec 2024 08:20:03 -0800
Subject: [PATCH] Add spec.find_site().

PiperOrigin-RevId: 702733068
Change-Id: I4cc7fd2b9b59185a4458fdac8537f8287af1f295
---
 python/mujoco/specs.cc      | 6 ++++++
 python/mujoco/specs_test.py | 1 +
 2 files changed, 7 insertions(+)

diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc
index 9589229d..047002f9 100644
--- a/python/mujoco/specs.cc
+++ b/python/mujoco/specs.cc
@@ -344,6 +344,12 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_findFrame(self.ptr, name.c_str());
       },
       py::return_value_policy::reference_internal);
+  mjSpec.def(
+      "find_site",
+      [](MjSpec& self, std::string& name) -> raw::MjsSite* {
+        return mjs_asSite(mjs_findElement(self.ptr, mjOBJ_SITE, name.c_str()));
+      },
+      py::return_value_policy::reference_internal);
   mjSpec.def(
       "find_default",
       [](MjSpec& self, std::string& classname) -> const raw::MjsDefault* {
diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index afbcc576..a7133e77 100644
--- a/python/mujoco/specs_test.py
+++ b/python/mujoco/specs_test.py
@@ -723,6 +723,7 @@ class SpecsTest(absltest.TestCase):
     self.assertIsNotNone(head)
     site = head.first_site()
     self.assertIsNotNone(site)
+    self.assertEqual(site, spec.find_site('head'))
 
     site.delete()
     spec.sensors[-1].delete()