diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index 6507b764..2d13667b 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -52,6 +52,15 @@ SPECS = [ ('mjsExclude', 'Spec', False, 'excludes', 'mjOBJ_EXCLUDE'), ('mjsPlugin', 'Spec', False, 'plugins', 'mjOBJ_PLUGIN'), ] +SPECS_ADD = SPECS + [ + ('mjsBody', 'Frame', True, 'bodies', 'mjOBJ_BODY'), + ('mjsSite', 'Frame', True, 'sites', 'mjOBJ_SITE'), + ('mjsGeom', 'Frame', True, 'geoms', 'mjOBJ_GEOM'), + ('mjsJoint', 'Frame', True, 'joints', 'mjOBJ_JOINT'), + ('mjsCamera', 'Frame', True, 'cameras', 'mjOBJ_CAMERA'), + ('mjsFrame', 'Frame', True, 'frames', 'mjOBJ_FRAME'), + ('mjsLight', 'Frame', True, 'lights', 'mjOBJ_LIGHT'), +] # pylint: enable=bad-whitespace @@ -301,7 +310,7 @@ def generate() -> None: def generate_add() -> None: """Generate add constructors with optional keyword arguments.""" - for key, parent, default, listname, objtype in SPECS: + for key, parent, default, listname, objtype in SPECS_ADD: def _field(f: ast_nodes.StructFieldDecl): if f.type == ast_nodes.PointerType( @@ -390,6 +399,22 @@ def generate_add() -> None: raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{ auto out = mjs_add{elem}(&self, default_); """ + elif parent == 'Frame': + if key == 'mjsFrame': + code = f""" + {'mjs' + parent}.def("add_{elemlower}", []({'raw::Mjs' + parent}& self, + raw::MjsFrame* parentframe_, py::kwargs kwargs) -> raw::{titlecase}* {{ + raw::MjsBody* body = mjs_getParent(self.element); + auto out = mjs_add{elem}(body, &self); + """ + else: + code = f""" + {'mjs' + parent}.def("add_{elemlower}", []({'raw::Mjs' + parent}& self, + raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{ + raw::MjsBody* body = mjs_getParent(self.element); + auto out = mjs_add{elem}(body, default_); + mjs_setFrame(out->element, &self); + """ else: raise NotImplementedError(f'{parent} parent is not implement.') diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 305a4c4b..2e1b6c4e 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -318,6 +318,14 @@ class SpecsTest(absltest.TestCase): np.testing.assert_array_equal(frameb1.pos, frameb0.pos) np.testing.assert_array_equal(frameb1.quat, frameb0.quat) + # Add frame in frame. + framec0 = frameb0.add_frame(name='framec', pos=[7, 8, 9], quat=[0, 0, 0, 1]) + self.assertEqual(framec0.name, 'framec') + self.assertEqual(framec0.parent, frameb0.parent) + self.assertEqual(framec0.frame, frameb0) + np.testing.assert_array_equal(framec0.pos, [7, 8, 9]) + np.testing.assert_array_equal(framec0.quat, [0, 0, 0, 1]) + # Add joint. joint = body.add_joint(type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 1, 0]) self.assertEqual(joint.type, mujoco.mjtJoint.mjJNT_HINGE) @@ -340,6 +348,12 @@ class SpecsTest(absltest.TestCase): light = body.add_light(attenuation=[1, 2, 3]) np.testing.assert_array_equal(light.attenuation, [1, 2, 3]) + # Add light in a frame. + light_in_frame = framea0.add_light(cutoff=10) + self.assertEqual(light_in_frame.cutoff, 10) + self.assertEqual(light_in_frame.parent, framea0.parent) + self.assertEqual(light_in_frame.frame, framea0) + # Invalid input for valid keyword argument. with self.assertRaises(ValueError) as cm: body.add_geom(pos='pos')