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'" ):