Add from_zip to parser.

PiperOrigin-RevId: 809703253
Change-Id: I12b77a7d7c539419c2e3ee4fce50ed04e604a3cc

Author: Adrian Collister

dm_control/mjcf/__init__.py

@@ -29,6 +29,7 @@
 from dm_control.mjcf.parser import from_file
 from dm_control.mjcf.parser import from_path
 from dm_control.mjcf.parser import from_xml_string
+from dm_control.mjcf.parser import from_zip
 
 from dm_control.mjcf.physics import Physics
 

dm_control/mjcf/parser.py

@@ -15,8 +15,10 @@
 
 """Functions for parsing XML into an MJCF object model."""
 
+import io
 import os
 import sys
+import zipfile
 
 from dm_control.mjcf import constants
 from dm_control.mjcf import debugging
@@ -107,6 +109,60 @@ def from_path(path, escape_separators=False, resolve_references=True,
                 assets=assets)
 
 
+def from_zip(path, model_file='model.xml', escape_separators=False,
+             resolve_references=True):
+  """Parses a zipped XML file into an MJCF object model.
+
+  Args:
+    path: A path to a zip file containing an MJCF model and its assets.
+    model_file: If the zip contains multiple XML files, specify the name of the
+      main model file. Ignored if the zip only contains one XML file.
+    escape_separators: (optional) A boolean, whether to replace '/' characters
+      in element identifiers. If `False`, any '/' present in the XML causes a
+      ValueError to be raised.
+    resolve_references: (optional) A boolean indicating whether the parser
+      should attempt to resolve reference attributes to a corresponding element.
+
+  Returns:
+    An `mjcf.RootElement`.
+
+  Raises:
+    ValueError: If:
+      - the path does not point to a zip file
+      - the zip file contains no XML files
+      - the zip file contains more than one XML file and none of them have the
+        name specified in `model_file`.
+  """
+  contents = resources.GetResource(path)
+  if not zipfile.is_zipfile(io.BytesIO(contents)):
+    raise ValueError(f'File {path} is not a zip file.')
+  with zipfile.ZipFile(io.BytesIO(contents), 'r') as zf:
+    xml_files = [f for f in zf.namelist() if f.endswith('.xml')]
+    if not xml_files:
+      raise ValueError(f'No XML file found in {path}.')
+    elif len(xml_files) > 1:
+      model_files = [f for f in xml_files if f == model_file]
+      if not model_files:
+        raise ValueError(
+            f'Multiple XML files found in {path}, but none named {model_file}.'
+        )
+      xml_path = model_files[0]
+    else:
+      xml_path = xml_files[0]
+    xml_string = zf.read(xml_path)
+
+    model_dir = os.path.dirname(xml_path)
+    assets = {
+        os.path.relpath(name, model_dir): zf.read(name)
+        for name in zf.namelist()
+        if not (name.endswith(os.path.sep) or name == xml_path)
+    }
+
+  xml_root = etree.fromstring(xml_string)
+  return _parse(xml_root, escape_separators,
+                resolve_references=resolve_references, assets=assets)
+
+
 def _parse(xml_root, escape_separators=False,
            model_dir='', resolve_references=True, assets=None):
   """Parses a complete MJCF model from an XML.

dm_control/mjcf/test_assets/model_with_assetdir.zip

@@ -25,9 +25,10 @@
 _ARENA_XML = os.path.join(ASSETS_DIR, 'arena.xml')
 _LEGO_BRICK_XML = os.path.join(ASSETS_DIR, 'lego_brick.xml')
 _ROBOT_XML = os.path.join(ASSETS_DIR, 'robot_arm.xml')
+_ZIPPED_MODEL = os.path.join(ASSETS_DIR, 'model_with_assetdir.zip')
 
 
-def validate(xml_string):
+def validate(xml_string, assets=None):
   """Validates that an XML string is a valid MJCF.
 
   Validation is performed by constructing Mujoco model from the string.
@@ -36,9 +37,10 @@ def validate(xml_string):
 
   Args:
     xml_string: XML string to validate
+    assets: Optional dict of assets to use for the model.
   """
 
-  mjmodel = wrapper.MjModel.from_xml_string(xml_string)
+  mjmodel = wrapper.MjModel.from_xml_string(xml_string, assets)
   wrapper.MjData(mjmodel)
 
 
@@ -61,6 +63,10 @@ def testXmlAttach(self):
     # validate
     validate(arena.to_xml_string())
 
+  def testXmlFromZip(self):
+    model = parser.from_zip(_ZIPPED_MODEL)
+    validate(model.to_xml_string(), model.get_assets())
+
 
 if __name__ == '__main__':
   absltest.main()

dm_control/mjcf/xml_validation_test.py

@@ -222,6 +222,7 @@ def is_excluded(s):
                 '*.xml',
                 '*.textproto',
                 '*.h5',
+                '*.zip',
             ],
             excludes=[
                 '*/dog_assets/extras/*',