Add mujoco.MjSpec.to_zip()

PiperOrigin-RevId: 719003836
Change-Id: I310ac96b0080b3313c325cee76d3bf7f81fa9df0

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

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'