diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py index 1cb446f7..2d8d22ad 100644 --- a/python/mujoco/__init__.py +++ b/python/mujoco/__init__.py @@ -124,15 +124,25 @@ def from_zip(file: Union[str, IO[bytes]]) -> _specs.MjSpec: if not zipfile.is_zipfile(file): raise ValueError(f'File {file} is not a zip file.') with zipfile.ZipFile(file, 'r') as zip_file: + xml_dir = None 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() + xml_dir = os.path.dirname(zip_info.filename) else: assets[zip_info.filename] = f.read() + if not xml_string: raise ValueError('No XML file found in zip file.') + + relative_assets = {} + for key, value in assets.items(): + new_key = os.path.relpath(key, xml_dir) + relative_assets[new_key] = value + assets = relative_assets + return _specs.MjSpec.from_string(xml_string, assets=assets) diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index ee99ed25..e2610f34 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -1615,5 +1615,25 @@ class SpecsTest(absltest.TestCase): self.assertEqual(wrap_joint2.type, mujoco.mjtWrap.mjWRAP_JOINT) self.assertEqual(wrap_joint2.coef, 2.0) + def test_from_zip(self): + """Tests that the assets are correctly parsed from a zip file.""" + model_path_root = ( + epath.resource_path("mujoco") / "testdata" / "MJCF_Root.zip" + ) + model_path_no_root = ( + epath.resource_path("mujoco") / "testdata" / "MJCF_NoRoot.zip" + ) + filenames = [model_path_root.as_posix(), model_path_no_root.as_posix()] + + for filename in filenames: + with self.subTest(filename): + spec = mujoco.MjSpec.from_zip(filename) + spec.compile() + assets = spec.assets + xml_string = spec.to_xml() + string_spec = mujoco.MjSpec.from_string(xml_string, assets=assets) + string_spec.compile() + self.assertEqual(spec.to_xml(), string_spec.to_xml()) + if __name__ == '__main__': absltest.main() diff --git a/python/mujoco/testdata/MJCF_NoRoot.zip b/python/mujoco/testdata/MJCF_NoRoot.zip new file mode 100644 index 00000000..215b9a68 Binary files /dev/null and b/python/mujoco/testdata/MJCF_NoRoot.zip differ diff --git a/python/mujoco/testdata/MJCF_Root.zip b/python/mujoco/testdata/MJCF_Root.zip new file mode 100644 index 00000000..9bf4d0ce Binary files /dev/null and b/python/mujoco/testdata/MJCF_Root.zip differ diff --git a/python/pyproject.toml b/python/pyproject.toml index f42aa7f9..d79683da 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -57,6 +57,7 @@ mujoco = [ "testdata/*.xml", "testdata/*.msh", "testdata/*.usda", + "testdata/*.zip", ] [project.optional-dependencies]