diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index ca8e8bdc..74ded675 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -316,13 +316,24 @@ def generate_add() -> None: if f.type == ast_nodes.PointerType( inner_type=ast_nodes.ValueType(name='mjsElement') ): - return '', '', '' + return '', '', '', '', '', '' elif f.type == ast_nodes.ValueType(name='mjsPlugin'): - return f'set_plugin(out->{f.name});', 'plugin', f.name + return ( + f'set_plugin(out->{f.name}, plugin);', + 'plugin', + f.name, + 'MjsPlugin', + 'std::optional& plugin', + 'py::arg("plugin") = py::none()', + ) elif f.type == ast_nodes.ValueType(name='mjsOrientation'): return ( ( f'set_orientation(out->{f.name},' + f' {"iaxisangle" if f.name == "ialt" else "axisangle"},' + f' {"ixyaxes" if f.name == "ialt" else "xyaxes"},' + f' {"izaxis" if f.name == "ialt" else "zaxis"},' + f' {"ieuler" if f.name == "ialt" else "euler"},' f' "{"iaxisangle" if f.name == "ialt" else "axisangle"}",' f' "{"ixyaxes" if f.name == "ialt" else "xyaxes"}",' f' "{"izaxis" if f.name == "ialt" else "zaxis"}",' @@ -332,85 +343,216 @@ def generate_add() -> None: ['iaxisangle', 'ixyaxes', 'izaxis', 'ieuler'] if f.name == 'ialt' else ['axisangle', 'xyaxes', 'zaxis', 'euler'], + 'list[float]', + [ + f'std::optional>& {n}' + for n in ( + ['iaxisangle', 'ixyaxes', 'izaxis', 'ieuler'] + if f.name == 'ialt' + else ['axisangle', 'xyaxes', 'zaxis', 'euler'] + ) + ], + [ + f'py::arg("{n}") = py::none()' + for n in ( + ['iaxisangle', 'ixyaxes', 'izaxis', 'ieuler'] + if f.name == 'ialt' + else ['axisangle', 'xyaxes', 'zaxis', 'euler'] + ) + ], ) elif f.type == ast_nodes.PointerType( inner_type=ast_nodes.ValueType(name='mjString') ): - return f'set_string("{f.name}", out->{f.name});', 'string', f.name + return ( + f'set_string(out->{f.name}, {f.name});', + 'string', + f.name, + 'str', + f'std::optional& {f.name}', + f'py::arg("{f.name}") = py::none()', + ) + elif f.type == ast_nodes.PointerType( + inner_type=ast_nodes.ValueType(name='mjStringVec') + ): + return ( + f'set_str_vec(out->{f.name}, {f.name});', + 'str_vec', + f.name, + 'list[str]', + f'std::optional>& {f.name}', + f'py::arg("{f.name}") = py::none()', + ) + + # Handle other vector types + inner_name = ( + f.type.inner_type.name + if isinstance(f.type, ast_nodes.PointerType) + and isinstance(f.type.inner_type, ast_nodes.ValueType) + else '' + ) + + if inner_name in ('mjIntVec', 'mjByteVec'): + return ( + f'set_int_vec(out->{f.name}, {f.name});', + 'int_vec', + f.name, + 'list[int]', + f'std::optional>& {f.name}', + f'py::arg("{f.name}") = py::none()', + ) + elif inner_name in ('mjFloatVec', 'mjDoubleVec'): + return ( + f'set_vec(out->{f.name}, {f.name});', + 'vec', + f.name, + 'list[float]', + f'std::optional>& {f.name}', + f'py::arg("{f.name}") = py::none()', + ) + elif inner_name == 'mjIntVecVec': + return ( + f'set_int_vec_vec(out->{f.name}, {f.name});', + 'int_vec_vec', + f.name, + 'list[list[int]]', + f'std::optional>>& {f.name}', + f'py::arg("{f.name}") = py::none()', + ) + elif inner_name == 'mjFloatVecVec': + return ( + f'set_float_vec_vec(out->{f.name}, {f.name});', + 'float_vec_vec', + f.name, + 'list[list[float]]', + f'std::optional>>& {f.name}', + f'py::arg("{f.name}") = py::none()', + ) + elif isinstance(f.type, ast_nodes.PointerType): - return f'set_vec("{f.name}", out->{f.name});', 'vec', f.name + return ( + f'set_vec(out->{f.name}, {f.name});', + 'vec', + f.name, + 'list[float]', + f'std::optional>& {f.name}', + f'py::arg("{f.name}") = py::none()', + ) elif isinstance(f.type, ast_nodes.ArrayType): return ( - f'set_array("{f.name}", out->{f.name}, {f.type.extents[0]});', + ( + f'set_array(out->{f.name}, {f.name}, {f.type.extents[0]},' + f' "{f.name}");' + ), 'array', f.name, + 'list[float]', + f'std::optional>& {f.name}', + f'py::arg("{f.name}") = py::none()', ) elif isinstance(f.type, ast_nodes.ValueType): - return f'set_value("{f.name}", out->{f.name});', 'value', f.name + type_name = 'float' + cpp_type = 'double' + if f.type.name in ('int', 'mjtByte') or f.type.name.startswith('mjt'): + type_name = 'int' + cpp_type = 'int' + return ( + f'set_value(out->{f.name}, {f.name});', + 'value', + f.name, + type_name, + f'std::optional<{cpp_type}>& {f.name}', + f'py::arg("{f.name}") = py::none()', + ) else: - return '', '', '' + return '', '', '', '', '', '' if key == 'mjsPlugin': code_field = '' set_types = [] names = [] + types = [] + cpp_args = [] + py_args = [] else: - code_field = 'set_name("name", out->element);' + code_field = 'set_name(out->element, name);' set_types = ['name'] names = ['name'] + types = ['str'] + cpp_args = ['std::optional& name'] + py_args = ['py::arg("name") = py::none()'] + for field in structs.STRUCTS[key].fields: - line, set_type, name = _field(field) + line, set_type, name, type_name, cpp_arg, py_arg = _field(field) if line: code_field = code_field + '\n ' + line set_types.append(set_type) if set_type == 'orientation': names.extend(name) + types.extend([type_name] * len(name)) + cpp_args.extend(cpp_arg) + py_args.extend(py_arg) else: names.append(name) + types.append(type_name) + cpp_args.append(cpp_arg) + py_args.append(py_arg) # assemble elem = key.removeprefix('mjs') elemlower = elem.lower() titlecase = 'Mjs' + elem + docstring = f'Add {elemlower} to spec.\n\n Args:\n' + for i, name in enumerate(names): + docstring += f'{'\n' if i > 0 else ''} {name}: {types[i]}' + + # functions arguments + args = ', '.join(cpp_args) + if args: + args = ', ' + args + + # py::arg definitions + pyargs = ', '.join(py_args) + # function definition and call to mjs_add_ if parent == 'Spec': if default: code = f""" {'mj' + parent}.def("add_{elemlower}", []({'Mj' + parent}& self, - raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{ + raw::MjsDefault* default_{args}) -> raw::{titlecase}* {{ auto out = mjs_add{elem}(self.ptr, default_); """ else: code = f""" - {'mj' + parent}.def("add_{elemlower}", []({'Mj' + parent}& self, py::kwargs kwargs) -> raw::{titlecase}* {{ + {'mj' + parent}.def("add_{elemlower}", []({'Mj' + parent}& self{args}) -> raw::{titlecase}* {{ auto out = mjs_add{elem}(self.ptr); """ elif parent == 'Body': if key == 'mjsFrame': code = f""" {'mjs' + parent}.def("add_{elemlower}", []({'raw::Mjs' + parent}& self, - raw::MjsFrame* parentframe_, py::kwargs kwargs) -> raw::{titlecase}* {{ + raw::MjsFrame* parentframe_{args}) -> raw::{titlecase}* {{ auto out = mjs_add{elem}(&self, parentframe_); """ else: code = f""" {'mjs' + parent}.def("add_{elemlower}", []({'raw::Mjs' + parent}& self, - raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{ + raw::MjsDefault* default_{args}) -> 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::MjsFrame* parentframe_{args}) -> 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::MjsDefault* default_{args}) -> raw::{titlecase}* {{ raw::MjsBody* body = mjs_getParent(self.element); auto out = mjs_add{elem}(body, default_); mjs_setFrame(out->element, &self); @@ -418,199 +560,184 @@ def generate_add() -> None: else: raise NotImplementedError(f'{parent} parent is not implement.') - # check for valid kwargs - code += '\n std::set valid_kwargs = {' - valid_kwargs = '' - for i, name in enumerate(names): - valid_kwargs += f'"{name}"' - if i != len(names) - 1: - valid_kwargs += ', ' - code += valid_kwargs + '};' - - code += f"""\n - py::dict kwarg_dict = kwargs; - for (auto item: kwarg_dict) {{ - std::string key = py::str(item.first); - if (valid_kwargs.count(key) == 0) {{ - throw pybind11::type_error("Invalid " - + key - + " keyword argument. Valid options are: {", ".join(names)}."); - }} - }} - """ - # include helper functions if set_types: for t in set(set_types): if t == 'orientation': code += """\n - auto set_orientation = [&kwargs](raw::MjsOrientation& orientation, - const char* axisangle, - const char* xyaxes, - const char* zaxis, - const char* euler) { + auto set_orientation = [](raw::MjsOrientation& orientation, + const std::optional>& axisangle, + const std::optional>& xyaxes, + const std::optional>& zaxis, + const std::optional>& euler, + const char* name_axisangle, + const char* name_xyaxes, + const char* name_zaxis, + const char* name_euler) { int nrepresentation = 0; - bool has_axisangle = kwargs.contains(axisangle); - nrepresentation += has_axisangle; - bool has_xyaxes = kwargs.contains(xyaxes); - nrepresentation += has_xyaxes; - bool has_zaxis = kwargs.contains(zaxis); - nrepresentation += has_zaxis; - bool has_euler = kwargs.contains(euler); - nrepresentation += has_euler; + nrepresentation += axisangle.has_value(); + nrepresentation += xyaxes.has_value(); + nrepresentation += zaxis.has_value(); + nrepresentation += euler.has_value(); if (nrepresentation == 0) { return; } else if (nrepresentation > 1) { - throw pybind11::value_error("Only one of: " - + std::string(axisangle) + ", " - + std::string(xyaxes) + ", " - + std::string(zaxis) - + ", or" - + std::string(euler) - + " can be set."); + std::string msg = std::string("Only one of: ") + name_axisangle + ", " + name_xyaxes + ", " + name_zaxis + ", or " + name_euler + " can be set."; + throw pybind11::value_error(msg); } - auto set_array = [&kwargs](const char* str, double* des, int size) { - try { - std::vector array = kwargs[str].cast>(); - if (array.size() != size) { - throw pybind11::value_error(std::string(str) - + " should be a list/array of size " - + std::to_string(size) - + "."); - } - int idx = 0; - for (auto val : array) { - des[idx++] = val; - } - } catch (const py::cast_error &e) { - throw pybind11::value_error(std::string(str) - + " should be a list/array."); + auto set_array = [](const std::vector& array, double* des, int size, const char* name) { + if (array.size() != size) { + std::string msg = std::string(name) + " should be a list/array of size " + std::to_string(size) + "."; + throw pybind11::value_error(msg); + } + int idx = 0; + for (auto val : array) { + des[idx++] = val; } }; - - if (has_axisangle) { - set_array(axisangle, orientation.axisangle, 4); + if (axisangle.has_value()) { + set_array(axisangle.value(), orientation.axisangle, 4, name_axisangle); orientation.type = mjORIENTATION_AXISANGLE; - } else if (has_xyaxes) { - set_array(xyaxes, orientation.xyaxes, 6); + } else if (xyaxes.has_value()) { + set_array(xyaxes.value(), orientation.xyaxes, 6, name_xyaxes); orientation.type = mjORIENTATION_XYAXES; - } else if (has_zaxis) { - set_array(zaxis, orientation.zaxis, 3); + } else if (zaxis.has_value()) { + set_array(zaxis.value(), orientation.zaxis, 3, name_zaxis); orientation.type = mjORIENTATION_ZAXIS; - } else if (has_euler) { - set_array(euler, orientation.euler, 3); + } else if (euler.has_value()) { + set_array(euler.value(), orientation.euler, 3, name_euler); orientation.type = mjORIENTATION_EULER; } }; """ elif t == 'plugin': code += """\n - auto set_plugin = [&kwargs](raw::MjsPlugin& plugin) { - if (kwargs.contains("plugin")) { - std::optional input = kwargs["plugin"].cast(); - if (input.has_value()) { - try { - plugin.name = input->name; - } catch (const py::cast_error &e) { - throw pybind11::value_error("plugin.name should be a string."); - } - try { - plugin.plugin_name = input->plugin_name; - } catch (const py::cast_error &e) { - throw pybind11::value_error("plugin.instance_name should be a string."); - } - try { - plugin.active = input->active; - } catch (const py::cast_error &e) { - throw pybind11::value_error("plugin.active should be an mjtByte."); - } - try { - plugin.info = input->info; - } catch (const py::cast_error &e) { - throw pybind11::value_error("plugin.info should be a string."); - } - } + auto set_plugin = [](raw::MjsPlugin& plugin, const std::optional& input) { + if (input.has_value()) { + plugin.name = input->name; + plugin.plugin_name = input->plugin_name; + plugin.active = input->active; + plugin.info = input->info; } }; """ elif t == 'string': code += """\n - auto set_string = [&kwargs](const char* str, std::basic_string* des) { - if (kwargs.contains(str)) { - try { - *des = kwargs[str].cast(); - } catch (const py::cast_error &e) { - throw pybind11::value_error(std::string(str) + " should be a string."); + auto set_string = [](std::basic_string* des, const std::optional& str) { + if (str.has_value()) { + *des = str.value(); + } + }; + """ + elif t == 'str_vec': + code += """\n + auto set_str_vec = [](auto&& des, const std::optional>& vec) { + if (vec.has_value()) { + des->clear(); + des->reserve(vec->size()); + for (const auto& val : vec.value()) { + des->push_back(val); } } }; """ + + elif t == 'int_vec': + code += """\n + auto set_int_vec = [](auto&& des, const std::optional>& vec) { + if (vec.has_value()) { + using T = typename std::decay_t::value_type; + des->clear(); + des->reserve(vec->size()); + for (auto val : vec.value()) { + des->push_back(static_cast(val)); + } + } + }; + """ + elif t == 'int_vec_vec': + code += """\n + auto set_int_vec_vec = [](auto&& des, const std::optional>>& vec) { + if (vec.has_value()) { + des->clear(); + des->reserve(vec->size()); + for (const auto& inner : vec.value()) { + using InnerT = typename std::decay_t::value_type; + InnerT inner_res; + inner_res.reserve(inner.size()); + using ValT = typename InnerT::value_type; + for (auto val : inner) { + inner_res.push_back(static_cast(val)); + } + des->push_back(inner_res); + } + } + }; + """ + elif t == 'float_vec_vec': + code += """\n + auto set_float_vec_vec = [](auto&& des, const std::optional>>& vec) { + if (vec.has_value()) { + des->clear(); + des->reserve(vec->size()); + for (const auto& inner : vec.value()) { + using InnerT = typename std::decay_t::value_type; + InnerT inner_res; + inner_res.reserve(inner.size()); + using ValT = typename InnerT::value_type; + for (auto val : inner) { + inner_res.push_back(static_cast(val)); + } + des->push_back(inner_res); + } + } + }; + """ + elif t == 'vec': code += """\n - auto set_vec = [&kwargs](const char* str, auto&& des) { - if (kwargs.contains(str)) { - try { - using T = typename std::decay_t::value_type; - std::vector vec = kwargs[str].cast>(); - des->clear(); - des->reserve(vec.size()); - for (auto val : vec) { - des->push_back(val); - } - } catch (const py::cast_error &e) { - throw pybind11::value_error(std::string(str) + " has the wrong type."); + auto set_vec = [](auto&& des, const std::optional>& vec) { + if (vec.has_value()) { + using T = typename std::decay_t::value_type; + des->clear(); + des->reserve(vec->size()); + for (auto val : vec.value()) { + des->push_back(static_cast(val)); } } }; """ elif t == 'array': code += """\n - auto set_array = [&kwargs](const char* str, auto&& des, int size) { - if (kwargs.contains(str)) { - try { - using T = std::remove_pointer_t>; - std::vector array = kwargs[str].cast>(); - if (array.size() != size) { - throw pybind11::value_error(std::string(str) - + " should be a list/array of size " - + std::to_string(size) - + "."); - } - int idx = 0; - for (auto val : array) { - des[idx++] = val; - } - } catch (const py::cast_error &e) { - throw pybind11::value_error(std::string(str) + " should be a list/array."); + auto set_array = [](auto&& des, const std::optional>& array, int size, const char* name) { + if (array.has_value()) { + if (array->size() != size) { + std::string msg = std::string(name) + " should be a list/array of size " + std::to_string(size) + "."; + throw pybind11::value_error(msg); + } + int idx = 0; + for (auto val : array.value()) { + des[idx++] = val; } } }; """ elif t == 'value': code += """\n - auto set_value = [&kwargs](const char* str, auto&& des) { - if (kwargs.contains(str)) { - try { - using T = std::decay_t; - des = kwargs[str].cast(); - } catch (const py::cast_error &e) { - throw pybind11::value_error(std::string(str) + " is the wrong type."); - } + auto set_value = [](auto&& des, auto&& val) { + if (val.has_value()) { + using T = std::decay_t; + des = static_cast(val.value()); } }; """ elif t == 'name': code += """\n - auto set_name = [&kwargs](const char* str, raw::MjsElement* el) { - if (kwargs.contains(str)) { - try { - std::string name = kwargs[str].cast(); - mjs_setName(el, name.c_str()); - } catch (const py::cast_error &e) { - throw pybind11::value_error(std::string(str) + " should be a string."); - } + auto set_name = [](raw::MjsElement* el, const std::optional& name) { + if (name.has_value()) { + mjs_setName(el, name->c_str()); } }; """ @@ -621,7 +748,11 @@ def generate_add() -> None: code += f"""\n return out; }}, - {'py::arg_v("default", nullptr),' if default else ''} + {'py::arg_v("default", nullptr)' + (', ' if pyargs else '') if default else ''} + {pyargs}{', ' if pyargs else ''} + R"mydelimiter( + {docstring} + )mydelimiter", py::return_value_policy::reference_internal); """ diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 565aca66..f8ba454e 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -357,12 +357,8 @@ class SpecsTest(absltest.TestCase): self.assertEqual(light_in_frame.frame, framea0) # Invalid input for valid keyword argument. - with self.assertRaises(ValueError) as cm: + with self.assertRaises(TypeError): body.add_geom(pos='pos') - self.assertEqual( - str(cm.exception), - 'pos should be a list/array.', - ) with self.assertRaises(ValueError) as cm: body.add_geom(pos=[0, 1]) @@ -371,32 +367,15 @@ class SpecsTest(absltest.TestCase): 'pos should be a list/array of size 3.', ) - with self.assertRaises(ValueError) as cm: + with self.assertRaises(TypeError): body.add_geom(type='type') - self.assertEqual( - str(cm.exception), - 'type is the wrong type.', - ) - with self.assertRaises(ValueError) as cm: + with self.assertRaises(TypeError): body.add_geom(userdata='') - self.assertEqual( - str(cm.exception), - 'userdata has the wrong type.', - ) # Invalid keyword argument. - with self.assertRaises(TypeError) as cm: + with self.assertRaises(TypeError): body.add_geom(vel='vel') - self.assertEqual( - str(cm.exception), - 'Invalid vel keyword argument. Valid options are: name, type, pos,' - ' quat, axisangle, xyaxes, zaxis, euler, fromto, size, contype,' - ' conaffinity, condim, priority, friction, solmix, solref, solimp,' - ' margin, gap, mass, density, typeinertia, fluid_ellipsoid,' - ' fluid_coefs, material, rgba, group, hfieldname, meshname, fitscale,' - ' userdata, plugin, info.', - ) # Orientation keyword arguments. geom_axisangle = body.add_geom(axisangle=[1, 2, 3, 4]) @@ -464,13 +443,13 @@ class SpecsTest(absltest.TestCase): body.add_geom(axisangle=[1, 2, 3, 4], euler=[1, 2, 3]) self.assertEqual( str(cm.exception), - 'Only one of: axisangle, xyaxes, zaxis, oreuler can be set.', + 'Only one of: axisangle, xyaxes, zaxis, or euler can be set.', ) with self.assertRaises(ValueError) as cm: spec.worldbody.add_body(iaxisangle=[1, 2, 3, 4], ieuler=[1, 2, 3]) self.assertEqual( str(cm.exception), - 'Only one of: iaxisangle, ixyaxes, izaxis, orieuler can be set.', + 'Only one of: iaxisangle, ixyaxes, izaxis, or ieuler can be set.', ) def test_load_xml(self):