Include external assets in to_zip output

python/mujoco/__init__.py

@@ -17,6 +17,7 @@
 import ctypes
 import ctypes.util
 import os
+from pathlib import Path
 import platform
 import subprocess
 from typing import Any, IO, Union, Sequence
@@ -106,11 +107,43 @@ def to_zip(spec: _specs.MjSpec, file: Union[str, IO[bytes]]) -> None:
     file: The path to the file to save to or the file object to write to.
   """
   files_to_zip = spec.assets
+
+  strip_path = bool(spec.strippath)
+  model_dir = spec.modelfiledir
+
+  def _resolve_path(filename: str, compiler_dir: str) -> Path:
+    if strip_path:
+      filename = os.path.basename(filename)
+
+    path = Path(filename)
+
+    if path.is_absolute():
+      return path
+
+    if compiler_dir and Path(compiler_dir).is_absolute():
+      return Path(compiler_dir) / path
+
+    base_dir = Path(model_dir) if model_dir else Path.cwd()
+    if compiler_dir:
+      return base_dir / compiler_dir / path
+
+    return base_dir / path
+
+  def _add_files(elements, compiler_dir: str):
+    for element in elements:
+      if element.file and element.file not in files_to_zip:
+        file_path = _resolve_path(element.file, compiler_dir)
+        if file_path.exists():
+          files_to_zip[element.file] = file_path.read_bytes()
+
+  _add_files(spec.meshes, spec.compiler.meshdir)
+  _add_files(spec.textures, spec.compiler.texturedir)
+  _add_files(spec.skins, spec.compiler.meshdir)
+  _add_files(spec.hfields, spec.compiler.meshdir)
+
   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')
+    Path(file).parent.mkdir(parents=True, exist_ok=True)
   with zipfile.ZipFile(file, 'w') as zip_file:
     for filename, contents in files_to_zip.items():
       zip_info = zipfile.ZipInfo(filename)

python/mujoco/specs_test.py

@@ -18,6 +18,7 @@
 import inspect
 import math
 import os
+from pathlib import Path
 import textwrap
 import typing
 import zipfile  # pylint: disable=unused-import
@@ -1744,6 +1745,59 @@ def test_from_zip(self):
         string_spec.compile()
         self.assertEqual(spec.to_xml(), string_spec.to_xml())
 
+  def test_to_zip_includes_assets(self):
+    testdata_path = epath.resource_path("mujoco") / "testdata"
+    asset_filename = 'abdomen_1_body.msh'
+    texture_filename = 'checkerboard.png'
+
+    temp_dir = self.create_tempdir()
+
+    import shutil
+    shutil.copy(testdata_path / asset_filename, temp_dir.full_path)
+    shutil.copy(testdata_path / texture_filename, temp_dir.full_path)
+
+    cwd = os.getcwd()
+    self.addCleanup(os.chdir, cwd)
+    os.chdir(temp_dir.full_path)
+
+    spec = mujoco.MjSpec()
+    spec.modelname = 'test_model'
+    mesh = spec.add_mesh()
+    mesh.name = 'test_mesh'
+    mesh.file = asset_filename
+
+    texture = spec.add_texture()
+    texture.file = texture_filename
+    texture.type = mujoco.mjtTexture.mjTEXTURE_2D
+
+    zip_filename = 'model.zip'
+    spec.to_zip(zip_filename)
+
+    with zipfile.ZipFile(zip_filename, 'r') as z:
+      self.assertIn(asset_filename, z.namelist())
+      with open(asset_filename, 'rb') as f:
+        self.assertEqual(z.read(asset_filename), f.read())
+      self.assertIn(texture_filename, z.namelist())
+      with open(texture_filename, 'rb') as f:
+        self.assertEqual(z.read(texture_filename), f.read())
+      self.assertIn('test_model.xml', z.namelist())
+
+    spec1 = mujoco.MjSpec.from_zip(zip_filename)
+    model = spec1.compile()
+
+  def test_to_zip_includes_assets_from_references(self):
+    temp_dir = self.create_tempdir()
+    temp_dir_path = epath.Path(temp_dir.full_path)
+
+    zip_path = temp_dir_path / "cube.zip"
+    mesh_model_path = epath.resource_path("mujoco") / "testdata" / "msh.xml"
+
+    spec0 = mujoco.MjSpec.from_file(mesh_model_path.as_posix())
+    spec0.to_zip(zip_path.as_posix())
+
+    spec1 = mujoco.MjSpec.from_zip(zip_path.as_posix())
+    model = spec1.compile()
+
   def test_rangefinder_sensor(self):
     """Test rangefinder sensor with mjSpec, iterative model building."""
     # Raydata field enum values for dataspec bitfield