From d1b02556ef141f46ad45a0ab11cc8742b80d0e65 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 4 Dec 2024 07:30:33 -0800 Subject: [PATCH] Allow keyword arguments in attach. PiperOrigin-RevId: 702718597 Change-Id: I1f1a5e9c88c8e102d5cbe5984511e7923ec7be3e --- python/mujoco/specs.cc | 47 ++++++++++++++++++++++++------------- python/mujoco/specs_test.py | 6 ++--- 2 files changed, 34 insertions(+), 19 deletions(-) diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index f97ed6fd..9589229d 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -625,15 +625,19 @@ PYBIND11_MODULE(_specs, m) { py::return_value_policy::reference_internal); mjsBody.def( "attach_frame", - [](raw::MjsBody& self, raw::MjsFrame& frame, std::string& prefix, - std::string& suffix) -> raw::MjsFrame* { - auto new_frame = - mjs_attachFrame(&self, &frame, prefix.c_str(), suffix.c_str()); + [](raw::MjsBody& self, raw::MjsFrame& frame, + std::optional& prefix, + std::optional& suffix) -> raw::MjsFrame* { + const char* p = prefix.has_value() ? prefix.value().c_str() : ""; + const char* s = suffix.has_value() ? suffix.value().c_str() : ""; + auto new_frame = mjs_attachFrame(&self, &frame, p, s); if (!new_frame) { throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element))); } return new_frame; }, + py::arg("frame"), py::arg("prefix") = py::none(), + py::arg("suffix") = py::none(), py::return_value_policy::reference_internal); mjsBody.def( "to_frame", @@ -654,34 +658,41 @@ PYBIND11_MODULE(_specs, m) { }); mjsFrame.def( "attach_body", - [](raw::MjsFrame& self, raw::MjsBody& body, std::string& prefix, - std::string& suffix) -> raw::MjsBody* { - auto new_body = - mjs_attachBody(&self, &body, prefix.c_str(), suffix.c_str()); + [](raw::MjsFrame& self, raw::MjsBody& body, + std::optional& prefix, + std::optional& suffix) -> raw::MjsBody* { + const char* p = prefix.has_value() ? prefix.value().c_str() : ""; + const char* s = suffix.has_value() ? suffix.value().c_str() : ""; + auto new_body = mjs_attachBody(&self, &body, p, s); if (!new_body) { throw pybind11::value_error( mjs_getError(mjs_getSpec(self.element))); } return new_body; }, + py::arg("body"), py::arg("prefix") = py::none(), + py::arg("suffix") = py::none(), py::return_value_policy::reference_internal); mjsFrame.def( "attach", - [](raw::MjsFrame& self, MjSpec& spec, std::string& prefix, - std::string& suffix) -> raw::MjsFrame* { + [](raw::MjsFrame& self, MjSpec& spec, std::optional& prefix, + std::optional& suffix) -> raw::MjsFrame* { auto world = mjs_findBody(spec.ptr, "world"); if (!world) { throw pybind11::value_error( mjs_getError(mjs_getSpec(self.element))); } - auto attached_world = - mjs_attachBody(&self, world, prefix.c_str(), suffix.c_str()); + const char* p = prefix.has_value() ? prefix.value().c_str() : ""; + const char* s = suffix.has_value() ? suffix.value().c_str() : ""; + auto attached_world = mjs_attachBody(&self, world, p, s); if (!attached_world) { throw pybind11::value_error( mjs_getError(mjs_getSpec(self.element))); } return mjs_bodyToFrame(&attached_world); }, + py::arg("spec"), py::arg("prefix") = py::none(), + py::arg("suffix") = py::none(), py::return_value_policy::reference_internal); // ============================= MJSGEOM ===================================== @@ -730,16 +741,20 @@ PYBIND11_MODULE(_specs, m) { py::return_value_policy::reference_internal); mjsSite.def( "attach", - [](raw::MjsSite& self, raw::MjsBody& body, std::string& prefix, - std::string& suffix) -> raw::MjsBody* { - auto new_body = - mjs_attachToSite(&self, &body, prefix.c_str(), suffix.c_str()); + [](raw::MjsSite& self, raw::MjsBody& body, + std::optional& prefix, + std::optional& suffix) -> raw::MjsBody* { + const char* p = prefix.has_value() ? prefix.value().c_str() : ""; + const char* s = suffix.has_value() ? suffix.value().c_str() : ""; + auto new_body = mjs_attachToSite(&self, &body, p, s); if (!new_body) { throw pybind11::value_error( mjs_getError(mjs_getSpec(self.element))); } return new_body; }, + py::arg("body"), py::arg("prefix") = py::none(), + py::arg("suffix") = py::none(), py::return_value_policy::reference_internal); // ============================= MJSCAMERA =================================== diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 7d42d3d0..afbcc576 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -867,7 +867,7 @@ class SpecsTest(absltest.TestCase): parent.compiler.degree = not child.compiler.degree body = child.worldbody.add_body(euler=[90, 0, 0]) frame = parent.worldbody.add_frame(euler=[-mujoco.mjPI / 2, 0, 0]) - frame.attach_body(body, 'child-', '') + frame.attach_body(body, prefix='child-') model = parent.compile() np.testing.assert_almost_equal(model.body_quat[1], [1, 0, 0, 0]) @@ -876,7 +876,7 @@ class SpecsTest(absltest.TestCase): parent = mujoco.MjSpec() site = parent.worldbody.add_site(pos=[1, 2, 3]) body = child.worldbody.add_body() - self.assertIsNotNone(site.attach(body, '_', '')) + self.assertIsNotNone(site.attach(body, prefix='_')) model = parent.compile() np.testing.assert_array_equal(model.body_pos[1], [1, 2, 3]) @@ -892,7 +892,7 @@ class SpecsTest(absltest.TestCase): child.worldbody.add_camera(name='camera') parent = mujoco.MjSpec() frame = parent.worldbody.add_frame(name='frame') - frame.attach(child, 'child-', '') + frame.attach(child, prefix='child-') self.assertLen(child.cameras, 1) self.assertLen(parent.bodies, 1) self.assertLen(parent.frames, 2)