From f480a9e9853b8db8767fa9005c81e304085fc89e Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Fri, 14 Mar 2025 08:30:11 -0700
Subject: [PATCH] Add MjSpec.from_zip() constructor.

PiperOrigin-RevId: 736871864
Change-Id: Ib6cc189d5912a57e5c505c744a4a1d6ad61d1c84
---
 python/mujoco/__init__.py   | 31 +++++++++++++++++++++++++++++--
 python/mujoco/specs_test.py |  5 ++++-
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py
index ca77a1a0..de578935 100644
--- a/python/mujoco/__init__.py
+++ b/python/mujoco/__init__.py
@@ -91,7 +91,7 @@ MjStruct: TypeAlias = Union[
 
 
 def to_zip(spec: _specs.MjSpec, file: Union[str, IO[bytes]]) -> None:
-  """Converts a spec to a zip file.
+  """Converts an MjSpec to a zip file.
 
   Args:
     spec: The mjSpec to save to a file.
@@ -109,6 +109,33 @@ def to_zip(spec: _specs.MjSpec, file: Union[str, IO[bytes]]) -> None:
       zip_file.writestr(zip_info, contents)
 
 
+def from_zip(file: Union[str, IO[bytes]]) -> _specs.MjSpec:
+  """Reads a zip file and returns an MjSpec.
+
+  Args:
+    file: The path to the file to read from or the file object to read from.
+  Returns:
+    An MjSpec object.
+  """
+  assets = {}
+  xml_string = None
+  if isinstance(file, str):
+    file = open(file, 'rb')
+  if not zipfile.is_zipfile(file):
+    raise ValueError(f'File {file} is not a zip file.')
+  with zipfile.ZipFile(file, 'r') as zip_file:
+    for zip_info in zip_file.infolist():
+      if not zip_info.filename.endswith(os.path.sep):
+        with zip_file.open(zip_info.filename) as f:
+          if zip_info.filename.endswith('.xml'):
+            xml_string = f.read()
+          else:
+            assets[zip_info.filename] = f.read()
+  if not xml_string:
+    raise ValueError('No XML file found in zip file.')
+  return _specs.MjSpec.from_string(xml_string, assets=assets)
+
+
 class _MjBindModel:
   def __init__(self, elements: Sequence[Any]):
     self.elements = elements
@@ -166,7 +193,7 @@ def _bind_data(
   else:
     return data.bind_scalar(specs)
 
-
+_specs.MjSpec.from_zip = from_zip
 _specs.MjSpec.to_zip = to_zip
 _structs.MjData.bind = _bind_data
 _structs.MjModel.bind = _bind_model
diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index ec2c0fbb..83493481 100644
--- a/python/mujoco/specs_test.py
+++ b/python/mujoco/specs_test.py
@@ -758,11 +758,14 @@ class SpecsTest(absltest.TestCase):
 
     xml = """
     <mujoco model="test">
+      <compiler angle="radian"/>
+
       <asset>
         <mesh name="cube" file="cube.obj"/>
       </asset>
+
       <worldbody>
-        <geom mesh="cube"/>
+        <geom type="mesh" mesh="cube"/>
       </worldbody>
     </mujoco>
     """