From 010d919244dcd06427dfa9c436153ff2aae468f3 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Wed, 29 Oct 2025 02:54:08 -0700
Subject: [PATCH] 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   | 16 ++++++++++++++--
 python/mujoco/specs_test.py |  2 ++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py
index d551daa2..1cb446f7 100644
--- a/python/mujoco/__init__.py
+++ b/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]
diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index 7a76194a..ee99ed25 100644
--- a/python/mujoco/specs_test.py
+++ b/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'"
     ):