From e389872c7db149f9efb1f1501a963f6ed6f12421 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Fri, 23 May 2025 09:23:28 -0700
Subject: [PATCH] Add frame property to python bindings.

PiperOrigin-RevId: 762452877
Change-Id: I555cec650d04700f6718b4d77a866c821b90685c
---
 python/mujoco/specs.cc      | 43 +++++++++++++++++++++++++++++++++++++
 python/mujoco/specs_test.py |  9 ++++++++
 2 files changed, 52 insertions(+)

diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc
index c1bbc3eb..b32942a5 100644
--- a/python/mujoco/specs.cc
+++ b/python/mujoco/specs.cc
@@ -736,6 +736,13 @@ PYBIND11_MODULE(_specs, m) {
         return frame;
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def(
+    "frame",
+      [](raw::MjsBody* self) -> raw::MjsFrame* {
+        return mjs_getFrame(self->element);
+      },
+      py::return_value_policy::reference_internal);
+
 
   // ============================= MJSFRAME ====================================
   mjsFrame.def("delete", [](raw::MjsFrame& self) { mjs_delete(self.element); });
@@ -767,6 +774,12 @@ PYBIND11_MODULE(_specs, m) {
       py::arg("body"), py::arg("prefix") = py::none(),
       py::arg("suffix") = py::none(),
       py::return_value_policy::reference_internal);
+  mjsFrame.def_property_readonly(
+      "frame",
+      [](raw::MjsFrame& self) -> raw::MjsFrame* {
+        return mjs_getFrame(self.element);
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSGEOM =====================================
   mjsGeom.def("delete", [](raw::MjsGeom& self) { mjs_delete(self.element); });
@@ -789,6 +802,12 @@ PYBIND11_MODULE(_specs, m) {
       [](raw::MjsGeom& self, raw::MjsDefault& default_) -> void {
         mjs_setDefault(self.element, &default_);
       });
+  mjsGeom.def_property_readonly(
+      "frame",
+      [](raw::MjsGeom& self) -> raw::MjsFrame* {
+        return mjs_getFrame(self.element);
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSJOINT ====================================
   mjsJoint.def("delete", [](raw::MjsJoint& self) { mjs_delete(self.element); });
@@ -811,6 +830,12 @@ PYBIND11_MODULE(_specs, m) {
       [](raw::MjsJoint& self, raw::MjsDefault& default_) -> void {
         mjs_setDefault(self.element, &default_);
       });
+  mjsJoint.def_property_readonly(
+      "frame",
+      [](raw::MjsJoint& self) -> raw::MjsFrame* {
+        return mjs_getFrame(self.element);
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSSITE =====================================
   mjsSite.def("delete", [](raw::MjsSite& self) { mjs_delete(self.element); });
@@ -850,6 +875,12 @@ PYBIND11_MODULE(_specs, m) {
       py::arg("body"), py::arg("prefix") = py::none(),
       py::arg("suffix") = py::none(),
       py::return_value_policy::reference_internal);
+  mjsSite.def_property_readonly(
+      "frame",
+      [](raw::MjsSite& self) -> raw::MjsFrame* {
+        return mjs_getFrame(self.element);
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSCAMERA ===================================
   mjsCamera.def("delete",
@@ -873,6 +904,12 @@ PYBIND11_MODULE(_specs, m) {
       [](raw::MjsCamera& self, raw::MjsDefault& default_) -> void {
         mjs_setDefault(self.element, &default_);
       });
+  mjsCamera.def_property_readonly(
+      "frame",
+      [](raw::MjsCamera& self) -> raw::MjsFrame* {
+        return mjs_getFrame(self.element);
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSLIGHT ====================================
   mjsLight.def("delete", [](raw::MjsLight& self) { mjs_delete(self.element); });
@@ -895,6 +932,12 @@ PYBIND11_MODULE(_specs, m) {
       [](raw::MjsLight& self, raw::MjsDefault& default_) -> void {
         mjs_setDefault(self.element, &default_);
       });
+  mjsLight.def_property_readonly(
+      "frame",
+      [](raw::MjsLight& self) -> raw::MjsFrame* {
+        return mjs_getFrame(self.element);
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSMATERIAL =================================
   mjsMaterial.def("delete",
diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index ac702876..c35c6618 100644
--- a/python/mujoco/specs_test.py
+++ b/python/mujoco/specs_test.py
@@ -1069,6 +1069,15 @@ class SpecsTest(absltest.TestCase):
     frame = body.to_frame()
     np.testing.assert_array_equal(frame.pos, [1, 2, 3])
 
+  def test_get_frame(self):
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body()
+    frame = body.add_frame()
+    geom = body.add_geom()
+    geom.set_frame(frame)
+    self.assertIsNotNone(frame)
+    self.assertIs(geom.frame, frame)
+
   def test_attach_to_frame(self):
     parent = mujoco.MjSpec()
     parent.assets = {'cube.obj': 'cube_content'}