From 41bcf2d621c9548c70608b8281b8e2609ce550ee Mon Sep 17 00:00:00 2001 From: Nimrod Gileadi Date: Tue, 28 Jun 2022 05:22:38 -0700 Subject: [PATCH] Python bindings: Add __repr__ implementations to grouped views. PiperOrigin-RevId: 457704673 Change-Id: Idf0d2132ce0872e6ae3eef319b4c3119ab769cd0 --- python/mujoco/bindings_test.py | 21 ++++++++++++++++++++- python/mujoco/structs.cc | 10 ++++++---- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index ddabfc57..cfda1ce8 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -34,7 +34,7 @@ TEST_XML = r""" - + @@ -263,6 +263,25 @@ class MuJoCoBindingsTest(parameterized.TestCase): self.data.joint('myfree').cdof = 42 np.testing.assert_array_equal(self.data.cdof[dof_idx:dof_idx+6], [[42]*6]*6) + def test_named_indexing_repr_in_data(self): + expected_repr = '''<_MjDataGeomViews + xmat: array([0., 0., 0., 0., 0., 0., 0., 0., 0.]) + xpos: array([0., 0., 0.]) +>''' + self.assertEqual(expected_repr, repr(self.data.geom('mybox'))) + + def test_named_indexing_body_repr_in_data(self): + view_repr = repr(self.data.body('mybox')) + self.assertStartsWith(view_repr, '<_MjDataBodyViews') + self.assertIn('xpos: array([0., 0., 0.])', view_repr) + self.assertEndsWith(view_repr, '>') + + def test_named_indexing_repr_in_model(self): + view_repr = repr(self.model.geom('mybox')) + self.assertStartsWith(view_repr, '<_MjModelGeomViews') + self.assertIn('size: array([0.1, 0.1, 0.1])', view_repr) + self.assertEndsWith(view_repr, '>') + def test_addresses_differ_between_structs(self): model2 = mujoco.MjModel.from_xml_string(TEST_XML) data2 = mujoco.MjData(model2) diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index ebeb2786..03f0ee26 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -1499,11 +1499,12 @@ This is useful for example when the MJB is not available as a file on disk.)")); MJMODEL_VIEW_GROUPS_ALTNAMES #undef XGROUP -#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \ - { \ - using GroupedViews = MjModelGroupedViews; \ +#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \ + { \ + using GroupedViews = MjModelGroupedViews; \ py::class_ groupedViews(m, "_" #MjModelGroupedViews); \ - FIELD_XMACROS \ + FIELD_XMACROS \ + groupedViews.def("__repr__", StructRepr); \ } #define X(type, prefix, var, dim0, dim1) \ groupedViews.def_property( \ @@ -1777,6 +1778,7 @@ This is useful for example when the MJB is not available as a file on disk.)")); using GroupedViews = MjDataGroupedViews; \ py::class_ groupedViews(m, "_" #MjDataGroupedViews); \ FIELD_XMACROS \ + groupedViews.def("__repr__", StructRepr); \ } #define X(type, prefix, var, dim0, dim1) \ groupedViews.def_property( \