Merge pull request #2496 from jjyyxx:patch-1

PiperOrigin-RevId: 736457846
Change-Id: I3b4b1554dd15e9ca73abee5c31d4758f694f5b44

python/mujoco/bindings_test.py

@@ -377,6 +377,12 @@ class MuJoCoBindingsTest(parameterized.TestCase):
         4,
     )
 
+  def test_mjvisual_repr(self):
+    # Regression test for issue #2488.
+    vis_repr = repr(self.model.vis)
+    self.assertNotEmpty(vis_repr)
+    self.assertIn('MjVisual', vis_repr)
+
   def test_mjmodel_can_read_and_write_opt(self):
     self.assertEqual(self.model.opt.timestep, 0.002)
     np.testing.assert_array_equal(self.model.opt.gravity, [0, 0, -9.81])

python/mujoco/structs.cc

@@ -1415,7 +1415,29 @@ PYBIND11_MODULE(_structs, m) {
   mjVisual.def("__deepcopy__", [](const MjVisualWrapper& other, py::dict) {
     return MjVisualWrapper(other);
   });
-  DefineStructFunctions(mjVisual);
+  mjVisual.def("__eq__", StructsEqual<MjVisualWrapper>);
+  // Special __repr__ implementation for MjVisual, since:
+  // 1. Types under MjVisual confuse StructRepr;
+  // 2. StructRepr does not handle the indentation of nested structs well.
+  mjVisual.def("__repr__", [](py::object self) {
+    std::ostringstream result;
+    result << "<"
+           << self.attr("__class__").attr("__name__").cast<std::string_view>();
+
+#define X(type, var)          \
+    result << "\n  " #var ": "; \
+    StructReprImpl<type>(self.attr(#var), result, 2);
+
+    X(raw::MjVisualGlobal, global_)
+    X(raw::MjVisualQuality, quality)
+    X(MjVisualHeadlightWrapper, headlight)
+    X(raw::MjVisualMap, map)
+    X(raw::MjVisualScale, scale)
+    X(MjVisualRgbaWrapper, rgba)
+#undef X
+    result << "\n>";
+    return result.str();
+  });
 
   py::class_<raw::MjVisualGlobal> mjVisualGlobal(mjVisual, "Global");
   mjVisualGlobal.def("__copy__", [](const raw::MjVisualGlobal& other) {

python/mujoco/structs.h

@@ -1206,8 +1206,8 @@ bool StructsEqual(pybind11::object lhs, pybind11::object rhs) {
 
 // Returns a string representation of a struct like object.
 template <typename T>
-std::string StructRepr(pybind11::object self) {
-  std::ostringstream result;
+void StructReprImpl(pybind11::object self, std::ostringstream& result,
+                    int indent) {
   result << "<"
          << self.attr("__class__").attr("__name__").cast<std::string_view>();
   for (pybind11::handle f : Dir<T>()) {
@@ -1216,10 +1216,16 @@ std::string StructRepr(pybind11::object self) {
       continue;
     }
 
-    result << "\n  " << name << ": "
+    result << "\n" << std::string(indent + 2, ' ') << name << ": "
            << self.attr(f).attr("__repr__")().cast<std::string_view>();
   }
-  result << "\n>";
+  result << "\n" << std::string(indent, ' ') << ">";
+}
+
+template <typename T>
+std::string StructRepr(pybind11::object self) {
+  std::ostringstream result;
+  StructReprImpl<T>(self, result, 0);
   return result.str();
 }