mjSpec python bindings for add functions: accept size with 1, 2, or 3 parameters

PiperOrigin-RevId: 858494291
Change-Id: I8a02dcac1d2419fbc728607907425b401cf66da4

python/mujoco/codegen/generate_spec_bindings.py

@@ -463,6 +463,15 @@ def generate_add() -> None:
             f'py::arg("{f.name}") = py::none()',
         )
       elif isinstance(f.type, ast_nodes.ArrayType):
+        if f.name == 'size' and f.type.extents[0] == 3:
+          return (
+              f'set_array_size(out->{f.name}, {f.name});',
+              'array_size',
+              f.name,
+              'list[float]',
+              f'std::optional<std::vector<double>>& {f.name}',
+              f'py::arg("{f.name}") = py::none()',
+          )
         return (
             (
                 f'set_array(out->{f.name}, {f.name}, {f.type.extents[0]},'
@@ -750,6 +759,20 @@ def generate_add() -> None:
             }
           };
           """
+        elif t == 'array_size':
+          code += """\n
+          auto set_array_size = [](auto&& des, const std::optional<std::vector<double>>& array) {
+            if (array.has_value()) {
+              if (array->size() < 1 || array->size() > 3) {
+                std::string msg = "size should be a list/array of size 1, 2, or 3.";
+                throw pybind11::value_error(msg);
+              }
+              for (int i = 0; i < 3; i++) {
+                des[i] = (i < array->size()) ? array->at(i) : 0;
+              }
+            }
+          };
+          """
         elif t == 'value':
           code += """\n
           auto set_value = [](auto&& des, auto&& val) {

python/mujoco/specs_test.py

@@ -452,6 +452,42 @@ class SpecsTest(absltest.TestCase):
         'Only one of: iaxisangle, ixyaxes, izaxis, or ieuler can be set.',
     )
 
+  def test_size_kwarg_variable_length(self):
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body()
+
+    geom_size1 = body.add_geom(size=[0.5])
+    np.testing.assert_array_equal(geom_size1.size, [0.5, 0, 0])
+
+    geom_size2 = body.add_geom(size=[0.5, 0.3])
+    np.testing.assert_array_equal(geom_size2.size, [0.5, 0.3, 0])
+
+    geom_size3 = body.add_geom(size=[0.5, 0.3, 0.1])
+    np.testing.assert_array_equal(geom_size3.size, [0.5, 0.3, 0.1])
+
+    site_size1 = body.add_site(size=[0.2])
+    np.testing.assert_array_equal(site_size1.size, [0.2, 0, 0])
+
+    site_size2 = body.add_site(size=[0.2, 0.1])
+    np.testing.assert_array_equal(site_size2.size, [0.2, 0.1, 0])
+
+    site_size3 = body.add_site(size=[0.2, 0.1, 0.05])
+    np.testing.assert_array_equal(site_size3.size, [0.2, 0.1, 0.05])
+
+    with self.assertRaises(ValueError) as cm:
+      body.add_geom(size=[])
+    self.assertEqual(
+        str(cm.exception),
+        'size should be a list/array of size 1, 2, or 3.',
+    )
+
+    with self.assertRaises(ValueError) as cm:
+      body.add_geom(size=[1, 2, 3, 4])
+    self.assertEqual(
+        str(cm.exception),
+        'size should be a list/array of size 1, 2, or 3.',
+    )
+
   def test_load_xml(self):
     file_path = epath.resource_path("mujoco") / "testdata" / "model.xml"
     filename = file_path.as_posix()