diff --git a/python/mujoco/render.cc b/python/mujoco/render.cc index b1e451ab..6c38348d 100644 --- a/python/mujoco/render.cc +++ b/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); diff --git a/python/mujoco/render_test.py b/python/mujoco/render_test.py index 5bb52615..a5ae4cfc 100644 --- a/python/mujoco/render_test.py +++ b/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() diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 481dd3b1..8eaa47a0 100644 --- a/python/mujoco/structs.cc +++ b/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 -py::object Dir() { - py::object type = py::type::of(); - auto dir = py::reinterpret_steal(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 -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()) { - auto name = f.cast(); - - 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 -std::string StructRepr(py::object self) { - std::ostringstream result; - result << "<" - << self.attr("__class__").attr("__name__").cast(); - for (py::handle f : Dir()) { - auto name = f.cast(); - if (name.empty() || name[0] == '_') { - continue; - } - - result << "\n " << name << ": " - << self.attr(f).attr("__repr__")().cast(); - } - result << "\n>"; - return result.str(); -} - -template -void DefineStructFunctions(py::class_ c) { - c.def("__eq__", StructsEqual); - c.def("__repr__", StructRepr); -} } // namespace // ==================== MJOPTION =============================================== diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h index 36c0e291..a7c5be42 100644 --- a/python/mujoco/structs.h +++ b/python/mujoco/structs.h @@ -973,6 +973,90 @@ static void DefinePyStr(pybind11::class_ 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 +pybind11::object Dir() { + pybind11::object type = pybind11::type::of(); + auto dir = + pybind11::reinterpret_steal(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 +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()) { + auto name = f.cast(); + + 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 +std::string StructRepr(pybind11::object self) { + std::ostringstream result; + result << "<" + << self.attr("__class__").attr("__name__").cast(); + for (pybind11::handle f : Dir()) { + auto name = f.cast(); + if (name.empty() || name[0] == '_') { + continue; + } + + result << "\n " << name << ": " + << self.attr(f).attr("__repr__")().cast(); + } + result << "\n>"; + return result.str(); +} + +template +void DefineStructFunctions(pybind11::class_ c) { + c.def("__eq__", StructsEqual); + c.def("__repr__", StructRepr); +} + } // namespace mujoco::python #endif // MUJOCO_PYTHON_STRUCTS_H_