diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py index 57b567b400..1637a2e636 100644 --- a/python/mujoco/__init__.py +++ b/python/mujoco/__init__.py @@ -19,7 +19,9 @@ import os import platform import subprocess +from typing import Union, IO import warnings +import zipfile # Extend the path to enable multiple directories to contribute to the same # package. Without this line, the `mujoco-mjx` package would not be able to @@ -53,10 +55,32 @@ from mujoco._functions import * from mujoco._render import * from mujoco._specs import * +from mujoco._specs import MjSpec from mujoco._structs import * from mujoco.gl_context import * from mujoco.renderer import Renderer + +def to_zip(spec: MjSpec, file: Union[str, IO[bytes]]) -> None: + """Converts a spec to a zip file. + + Args: + spec: The mjSpec to save to a file. + file: The path to the file to save to or the file object to write to. + """ + files_to_zip = spec.assets + files_to_zip[spec.modelname + '.xml'] = spec.to_xml() + if isinstance(file, str): + directory = os.path.dirname(file) + os.makedirs(directory, exist_ok=True) + file = open(file, 'wb') + with zipfile.ZipFile(file, 'w') as zip_file: + for filename, contents in files_to_zip.items(): + zip_info = zipfile.ZipInfo(os.path.join(spec.modelname, filename)) + zip_file.writestr(zip_info, contents) + +MjSpec.to_zip = to_zip + HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco') PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin') diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 4b53deedae..1302afc26e 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -15,7 +15,9 @@ """Tests for mjSpec bindings.""" import inspect +import os import textwrap +import zipfile from absl.testing import absltest from etils import epath @@ -735,6 +737,7 @@ def test_assets(self): v -1 -1 -1 v 1 -1 -1""" spec = mujoco.MjSpec() + spec.modelname = 'test' mesh = spec.add_mesh() mesh.name = 'cube' mesh.file = 'cube.obj'