Implement __repr__ and __eq__ on MjrRect in Python bindings.

PiperOrigin-RevId: 457457401
Change-Id: I0d395ff301fbcf775b4766ac81c830a625b3be63

python/mujoco/render.cc

@@ -155,6 +155,7 @@ PYBIND11_MODULE(_render, pymodule) {
   mjrRect.def("__deepcopy__", [](const raw::MjrRect& other, py::dict) {
     return raw::MjrRect(other);
   });
+  DefineStructFunctions(mjrRect);
 #define X(var) mjrRect.def_readwrite(#var, &raw::MjrRect::var)
   X(left);
   X(bottom);

python/mujoco/render_test.py

@@ -90,5 +90,12 @@ class MuJoCoRenderTest(absltest.TestCase):
     context.free()
     context.free()
 
+  def test_mjrrect_repr(self):
+    rect = mujoco.MjrRect(1, 2, 3, 4)
+    rect_repr = repr(rect)
+    self.assertIn('MjrRect', rect_repr)
+    self.assertIn('left: 1', rect_repr)
+
+
 if __name__ == '__main__':
   absltest.main()

python/mujoco/structs.cc

@@ -74,87 +74,6 @@ constexpr auto XArrayShapeImpl(const std::string_view dim1_str) {
   }
 }
 
-// An equals operator that works for comparing numpy arrays too.
-// Unlike other objects, the equality operator for numpy arrays returns a
-// numpy array.
-bool FieldsEqual(py::handle lhs, py::handle rhs, py::handle array_equal) {
-  // np.array_equal handles non-arrays.
-  return PyObject_IsTrue(array_equal(lhs, rhs).ptr());
-}
-
-// Returns an iterable object for iterating over attributes of T.
-template <typename T>
-py::object Dir() {
-  py::object type = py::type::of<T>();
-  auto dir = py::reinterpret_steal<py::object>(PyObject_Dir(type.ptr()));
-  if (PyErr_Occurred()) {
-    throw py::error_already_set();
-  }
-  return dir;
-}
-
-// Returns true if all public fields in lhs and rhs are equal.
-template <typename T>
-bool StructsEqual(py::object lhs, py::object rhs) {
-  // Equivalent to the following python code:
-  // if type(lhs) != type(rhs):
-  //   return False
-  // for field in dir(lhs):
-  //   if field.startswith("_"):
-  //     continue
-  //   # equal() handles equality of numpy arrays
-  //   if not equal(getattr(lhs, field, None), getattr(rhs, field, None)):
-  //     return False
-  //
-  // return True
-
-  auto np = py::module::import("numpy");
-  auto array_equal = np.attr("array_equal");
-
-  const py::handle lhs_t = py::type::handle_of(lhs);
-  const py::handle rhs_t = py::type::handle_of(rhs);
-  if (!lhs_t.is(rhs_t)) {
-    return false;
-  }
-  for (py::handle f : Dir<T>()) {
-    auto name = f.cast<std::string_view>();
-
-    if (name.empty() || name[0] == '_') {
-      continue;
-    }
-    py::object l = py::getattr(lhs, f, py::none());
-    py::object r = py::getattr(rhs, f, py::none());
-    if (!FieldsEqual(l, r, array_equal)) {
-      return false;
-    }
-  }
-  return true;
-}
-
-// Returns a string representation of a struct like object.
-template <typename T>
-std::string StructRepr(py::object self) {
-  std::ostringstream result;
-  result << "<"
-         << self.attr("__class__").attr("__name__").cast<std::string_view>();
-  for (py::handle f : Dir<T>()) {
-    auto name = f.cast<std::string_view>();
-    if (name.empty() || name[0] == '_') {
-      continue;
-    }
-
-    result << "\n  " << name << ": "
-           << self.attr(f).attr("__repr__")().cast<std::string_view>();
-  }
-  result << "\n>";
-  return result.str();
-}
-
-template <typename C, typename... O>
-void DefineStructFunctions(py::class_<C, O...> c) {
-  c.def("__eq__", StructsEqual<C>);
-  c.def("__repr__", StructRepr<C>);
-}
 }  // namespace
 
 // ==================== MJOPTION ===============================================

python/mujoco/structs.h

@@ -973,6 +973,90 @@ static void DefinePyStr(pybind11::class_<C, O...> c, const char* name,
       });
 }
 
+// An equals operator that works for comparing numpy arrays too.
+// Unlike other objects, the equality operator for numpy arrays returns a
+// numpy array.
+inline bool FieldsEqual(pybind11::handle lhs, pybind11::handle rhs,
+                        pybind11::handle array_equal) {
+  // np.array_equal handles non-arrays.
+  return PyObject_IsTrue(array_equal(lhs, rhs).ptr());
+}
+
+// Returns an iterable object for iterating over attributes of T.
+template <typename T>
+pybind11::object Dir() {
+  pybind11::object type = pybind11::type::of<T>();
+  auto dir =
+      pybind11::reinterpret_steal<pybind11::object>(PyObject_Dir(type.ptr()));
+  if (PyErr_Occurred()) {
+    throw pybind11::error_already_set();
+  }
+  return dir;
+}
+
+// Returns true if all public fields in lhs and rhs are equal.
+template <typename T>
+bool StructsEqual(pybind11::object lhs, pybind11::object rhs) {
+  // Equivalent to the following python code:
+  // if type(lhs) != type(rhs):
+  //   return False
+  // for field in dir(lhs):
+  //   if field.startswith("_"):
+  //     continue
+  //   # equal() handles equality of numpy arrays
+  //   if not equal(getattr(lhs, field, None), getattr(rhs, field, None)):
+  //     return False
+  //
+  // return True
+
+  auto np = pybind11::module::import("numpy");
+  auto array_equal = np.attr("array_equal");
+
+  const pybind11::handle lhs_t = pybind11::type::handle_of(lhs);
+  const pybind11::handle rhs_t = pybind11::type::handle_of(rhs);
+  if (!lhs_t.is(rhs_t)) {
+    return false;
+  }
+  for (pybind11::handle f : Dir<T>()) {
+    auto name = f.cast<std::string_view>();
+
+    if (name.empty() || name[0] == '_') {
+      continue;
+    }
+    pybind11::object l = pybind11::getattr(lhs, f, pybind11::none());
+    pybind11::object r = pybind11::getattr(rhs, f, pybind11::none());
+    if (!FieldsEqual(l, r, array_equal)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// Returns a string representation of a struct like object.
+template <typename T>
+std::string StructRepr(pybind11::object self) {
+  std::ostringstream result;
+  result << "<"
+         << self.attr("__class__").attr("__name__").cast<std::string_view>();
+  for (pybind11::handle f : Dir<T>()) {
+    auto name = f.cast<std::string_view>();
+    if (name.empty() || name[0] == '_') {
+      continue;
+    }
+
+    result << "\n  " << name << ": "
+           << self.attr(f).attr("__repr__")().cast<std::string_view>();
+  }
+  result << "\n>";
+  return result.str();
+}
+
+template <typename C, typename... O>
+void DefineStructFunctions(pybind11::class_<C, O...> c) {
+  c.def("__eq__", StructsEqual<C>);
+  c.def("__repr__", StructRepr<C>);
+}
+
 }  // namespace mujoco::python
 
 #endif  // MUJOCO_PYTHON_STRUCTS_H_