Yield ownership of char[], vector<string>, and vector<byte> to mjSpec in Python bindings.

PiperOrigin-RevId: 681522091
Change-Id: Icb3788d9ef6e914a307531e39ad8171c24f590af

python/mujoco/codegen/generate_spec_bindings.py

@@ -93,15 +93,15 @@ def _array_binding_code(
     return f"""\
   {classname}.def_property(
     "{varname}",
-    []({rawclassname}& self) -> py::array_t<char> {{
-      return py::array_t<char>({field.extents[0]}, self.{fullvarname});
+    []({rawclassname}& self) -> MjTypeVec<char> {{
+      return MjTypeVec<char>(self.{fullvarname}, {field.extents[0]});
     }},
     []({rawclassname}& self, py::object rhs) {{
       int i = 0;
       for (auto val : rhs) {{
         self.{fullvarname}[i++] = py::cast<char>(val);
       }}
-    }}, py::return_value_policy::reference_internal);"""
+    }}, py::return_value_policy::move);"""
   # all other array types
   return f"""\
   {classname}.def_property(
@@ -158,16 +158,13 @@ def _ptr_binding_code(
           self.{fullvarname}->push_back(py::cast<{vartype}>(val));
       }}
     }}, py::return_value_policy::reference_internal);"""
-  elif vartype == 'mjByteVec':  # C++ buffer -> Python list
+  elif vartype == 'mjByteVec':
     return f"""\
   {classname}.def_property(
     "{varname}",
-    []({rawclassname}& self) -> py::list {{
-        py::list list;
-        for (auto val : *self.{fullvarname}) {{
-          list.append(val);
-        }}
-        return list;
+    []({rawclassname}& self) -> MjTypeVec<std::byte> {{
+        return MjTypeVec<std::byte>(self.{fullvarname}->data(),
+                                    self.{fullvarname}->size());
       }},
     []({rawclassname}& self, py::object rhs) {{
         self.{fullvarname}->clear();
@@ -175,17 +172,14 @@ def _ptr_binding_code(
         for (auto val : rhs) {{
           self.{fullvarname}->push_back(py::cast<const std::byte>(val));
         }}
-    }}, py::return_value_policy::reference_internal);"""
-  elif vartype == 'mjStringVec':  # C++ vector of strings -> Python list
+    }}, py::return_value_policy::move);"""
+  elif vartype == 'mjStringVec':
     return f"""\
   {classname}.def_property(
     "{varname}",
-    []({rawclassname}& self) -> py::list {{
-        py::list list;
-        for (auto val : *self.{fullvarname}) {{
-          list.append(val);
-        }}
-        return list;
+    []({rawclassname}& self) -> MjTypeVec<std::string> {{
+        return MjTypeVec<std::string>(self.{fullvarname}->data(),
+                                      self.{fullvarname}->size());
       }},
     []({rawclassname}& self, py::object rhs) {{
         self.{fullvarname}->clear();
@@ -193,7 +187,7 @@ def _ptr_binding_code(
         for (auto val : rhs) {{
           self.{fullvarname}->push_back(py::cast<std::string>(val));
       }}
-    }}, py::return_value_policy::reference_internal);"""
+    }}, py::return_value_policy::move);"""
   elif 'VecVec' in vartype:  # C++ vector of vectors -> Python list of lists
     vartype = vartype.replace('mj', '').replace('VecVec', '').lower()
     return f"""\

python/mujoco/specs.cc

@@ -20,6 +20,7 @@
 #include <string>
 #include <string_view>  // IWYU pragma: keep
 #include <unordered_map>
+#include <utility>
 #include <vector>       // IWYU pragma: keep
 
 #include <Eigen/Core>
@@ -122,6 +123,38 @@ static raw::MjSpec* LoadSpecFileImpl(
   return spec;
 }
 
+template <typename T>
+struct MjTypeVec {
+  MjTypeVec(T* data, int size) : ptr(data), size(size) {}
+  T* ptr;
+  int size;
+};
+
+template <typename T>
+void DefineArray(py::module& m, const std::string& typestr) {
+  using Class = MjTypeVec<T>;
+  py::class_<Class>(m, typestr.c_str())
+      .def(py::init([](T* data, int size) { return Class(data, size); }))
+      .def("__getitem__",
+           [](Class& v, int i) {
+             if (i < 0 || i >= v.size) {
+               throw py::index_error("Index out of range.");
+             }
+             return v.ptr[i];
+           })
+      .def("__setitem__",
+           [](Class& v, int i, T c) {
+             if (i < 0 || i >= v.size) {
+               throw py::index_error("Index out of range.");
+             }
+             v.ptr[i] = std::move(c);
+           })
+      .def("__len__", [](Class& v) { return v.size; })
+      .def("__iter__", [](Class& v) {
+        return py::make_iterator(v.ptr, v.ptr + v.size);
+      }, py::keep_alive<0, 1>(), py::return_value_policy::reference_internal);
+};
+
 PYBIND11_MODULE(_specs, m) {
   auto structs_m = py::module::import("mujoco._structs");
   py::function mjmodel_from_spec_ptr =
@@ -161,6 +194,9 @@ PYBIND11_MODULE(_specs, m) {
   py::class_<raw::MjOption> mjOption(m, "MjOption");
   py::class_<raw::MjStatistic> mjStatistic(m, "MjStatistic");
   py::class_<raw::MjVisual> mjVisual(m, "MjVisual");
+  DefineArray<char>(m, "MjCharVec");
+  DefineArray<std::string>(m, "MjStringVec");
+  DefineArray<std::byte>(m, "MjByteVec");
 
   // ============================= MJSPEC =====================================
   mjSpec.def(py::init<>());

python/mujoco/specs_test.py

@@ -34,9 +34,28 @@ class SpecsTest(absltest.TestCase):
     spec = mujoco.MjSpec()
 
     # Check that euler sequence order is set correctly.
-    self.assertEqual(spec.eulerseq[0], ord('x'))
+    self.assertEqual(spec.eulerseq[0], 'x')
     spec.eulerseq = ['z', 'y', 'x']
-    self.assertEqual(spec.eulerseq[0], ord('z'))
+    self.assertEqual(spec.eulerseq[0], 'z')
+
+    # Change single elements of euler sequence.
+    spec.eulerseq[0] = 'y'
+    spec.eulerseq[1] = 'z'
+    self.assertEqual(spec.eulerseq[0], 'y')
+    self.assertEqual(spec.eulerseq[1], 'z')
+
+    # eulerseq is iterable
+    self.assertEqual('yzx', ''.join(spec.eulerseq))
+
+    # supports `len`
+    self.assertLen(spec.eulerseq, 3)
+
+    # field checks for out-of-bound access on read and on write
+    with self.assertRaises(IndexError):
+      spec.eulerseq[3] = 'x'
+
+    with self.assertRaises(IndexError):
+      spec.eulerseq[-1] = 'x'
 
     # Add a body, check that it has default orientation.
     body = spec.worldbody.add_body()
@@ -50,10 +69,14 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(body.name, 'baz')
 
     # Change the position of the body and read it back.
-    body.pos = [1, 2, 3]
-    np.testing.assert_array_equal(body.pos, [1, 2, 3])
+    body.pos = [4, 2, 3]
+    np.testing.assert_array_equal(body.pos, [4, 2, 3])
     self.assertEqual(body.pos.shape, (3,))
 
+    # Change single element of position.
+    body.pos[0] = 1
+    np.testing.assert_array_equal(body.pos, [1, 2, 3])
+
     # Change the orientation of the body and read it back.
     body.quat = [0, 1, 0, 0]
     np.testing.assert_array_equal(body.quat, [0, 1, 0, 0])
@@ -147,7 +170,7 @@ class SpecsTest(absltest.TestCase):
     # Add tuple.
     tuple_ = spec.add_tuple(objprm=[2.0, 3.0, 5.0], objname=['obj'])
     np.testing.assert_array_equal(tuple_.objprm, [2.0, 3.0, 5.0])
-    self.assertEqual(tuple_.objname, ['obj'])
+    self.assertEqual(tuple_.objname[0], 'obj')
 
     # Add flex.
     flex = spec.add_flex(friction=[1, 2, 3], texcoord=[1.0, 2.0, 3.0])
@@ -797,5 +820,31 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(model.stat.meansize, 0.06)
     self.assertEqual(model.vis.quality.shadowsize, 8192)
 
+  def test_assign_list_element(self):
+    spec = mujoco.MjSpec()
+    material = spec.add_material()
+    texture_index = mujoco.mjtTextureRole.mjTEXROLE_RGB
+
+    # Assign a string to a list element.
+    material.textures[texture_index] = 'texture_name'
+    self.assertEqual(material.textures[texture_index], 'texture_name')
+
+    # Assign a complete list
+    material.textures = ['', 'new_name', '', '', '']
+    self.assertEqual(material.textures[texture_index], 'new_name')
+
+    # textures is iterable
+    self.assertEqual('new_name', ''.join(material.textures))
+
+    # supports `len`
+    self.assertLen(material.textures, 5)
+
+    # field checks for out-of-bound access on read and on write
+    with self.assertRaises(IndexError):
+      material.textures[5] = 'x'
+
+    with self.assertRaises(IndexError):
+      material.textures[-1] = 'x'
+
 if __name__ == '__main__':
   absltest.main()