From 2546fcefa0e850265028a67ea1d88a83ff14cc10 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 23 Jan 2025 14:00:07 -0800
Subject: [PATCH] Add mujoco.MjSpec.to_zip()

PiperOrigin-RevId: 719003836
Change-Id: I310ac96b0080b3313c325cee76d3bf7f81fa9df0
---
 python/mujoco/__init__.py   | 24 ++++++++++++++++++++++++
 python/mujoco/specs_test.py |  3 +++
 2 files changed, 27 insertions(+)

diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py
index 57b567b4..1637a2e6 100644
--- a/python/mujoco/__init__.py
+++ b/python/mujoco/__init__.py
@@ -19,7 +19,9 @@ import ctypes.util
 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._errors import *
 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 4b53deed..1302afc2 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 @@ class SpecsTest(absltest.TestCase):
       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'