From 41968f1527944320f1df4970c93ac10c8ffbfeaa Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 5 Dec 2024 04:53:46 -0800
Subject: [PATCH] Support body.find_all("joint") in bindings.

PiperOrigin-RevId: 703071822
Change-Id: I6457ca0561c794447cc06aab29dcd0aac82c49a0
---
 python/mujoco/specs.cc      | 4 +++-
 python/mujoco/specs_test.py | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc
index 1a697124..c376ea35 100644
--- a/python/mujoco/specs.cc
+++ b/python/mujoco/specs.cc
@@ -493,6 +493,8 @@ PYBIND11_MODULE(_specs, m) {
           objtype = mjOBJ_GEOM;
         } else if (name == "site") {
           objtype = mjOBJ_SITE;
+        } else if (name == "joint") {
+          objtype = mjOBJ_JOINT;
         } else if (name == "light") {
           objtype = mjOBJ_LIGHT;
         } else if (name == "camera") {
@@ -500,7 +502,7 @@ PYBIND11_MODULE(_specs, m) {
         } else {
           throw pybind11::value_error(
               "body.find_all supports the types: body, frame, geom, site, "
-              "light, camera.");
+              "joint, light, camera.");
         }
         return FindAllImpl(self, objtype, true);
       },
diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index 91e260b0..e18db9b7 100644
--- a/python/mujoco/specs_test.py
+++ b/python/mujoco/specs_test.py
@@ -667,7 +667,7 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(
         str(cm.exception),
         'body.find_all supports the types: body, frame, geom, site,'
-        ' light, camera.',
+        ' joint, light, camera.',
     )
     body4 = spec.worldbody.find_all('body')[3]
     body4.name = 'body4_new'