Add __setattr__ to _MjBindModel and _MjBindData.

This allows setting attributes with bind using a list of mjSpec objects as arguments.

PiperOrigin-RevId: 825456920
Change-Id: Ie6af956feab0a805175d7aaecf26e071401fbf05

python/mujoco/__init__.py

@@ -137,8 +137,10 @@ def from_zip(file: Union[str, IO[bytes]]) -> _specs.MjSpec:
 
 
 class _MjBindModel:
+  """Wrapper for MjModel that allows binding multiple specs."""
+
   def __init__(self, elements: Sequence[Any]):
-    self.elements = elements
+    object.__setattr__(self, 'elements', elements)
 
   def __getattr__(self, key: str):
     items = []
@@ -146,10 +148,15 @@ class _MjBindModel:
       items.extend(getattr(e, key))
     return items
 
+  def __setattr__(self, key: str, value: Any):
+    raise AttributeError(f'Cannot set {key} on MjModel.')
+
 
 class _MjBindData:
+  """Wrapper for MjData that allows binding multiple specs."""
+
   def __init__(self, elements: Sequence[Any]):
-    self.elements = elements
+    object.__setattr__(self, 'elements', elements)
 
   def __getattr__(self, key: str):
     items = []
@@ -157,6 +164,11 @@ class _MjBindData:
       items.extend(getattr(e, key))
     return items
 
+  def __setattr__(self, key: str, value: Any):
+    value_it = iter(value)
+    for element in self.elements:
+      setattr(element, key, next(value_it))
+
 
 def _bind_model(
     model: _structs.MjModel, specs: Union[Sequence[MjStruct], MjStruct]

python/mujoco/specs_test.py

@@ -1266,6 +1266,8 @@ class SpecsTest(absltest.TestCase):
     np.testing.assert_array_equal(mj_model.bind(joints).qposadr, [7, 8])
     np.testing.assert_array_equal(mj_data.bind([]).qpos, [])
     np.testing.assert_array_equal(mj_model.bind([]).qposadr, [])
+    mj_data.bind(joints).qpos = np.array([1, 2])
+    np.testing.assert_array_equal(mj_data.bind(joints).qpos, [1, 2])
     with self.assertRaisesRegex(
         AttributeError, "object has no attribute 'invalid'"
     ):