Allow full model-relative paths to be used as keys in asset dicts.

Allowing keys in this form means the asset dict can be constructed without needing to first parse the XML to determine whether any special directories have been specified in the compiler options.

For example, given an asset in `assets/textures/texture.png` and an `assetdir` set to "assets", this asset previously had to be keyed in the assets dict by "textures/texture.png". After this CL, it's also permitted to use "assets/textures/texture.png".

PiperOrigin-RevId: 807216952
Change-Id: Ie1b8680eec447562675e4ca83aa15b1c4511461e

Author: Adrian Collister

dm_control/mjcf/attribute.py

@@ -518,39 +518,34 @@ def _get_asset_from_path(self, path):
     _, basename = os.path.split(path)
     filename, extension = os.path.splitext(basename)
 
-    assetdir = None
+    # Construct the path to the asset file (relative to the model directory),
+    # prefixed by the appropriate base path, if specified.
+    base_path = None
     if self._parent.namescope.has_identifier(
+        constants.BASEPATH, self._path_namespace
+    ):
+      base_path = self._parent.namescope.get(
+          constants.BASEPATH, self._path_namespace
+      )
+    elif self._path_namespace in (
+        constants.TEXTUREDIR_NAMESPACE,
+        constants.MESHDIR_NAMESPACE,
+    ) and self._parent.namescope.has_identifier(
         constants.BASEPATH, constants.ASSETDIR_NAMESPACE
     ):
-      assetdir = self._parent.namescope.get(
+      base_path = self._parent.namescope.get(
           constants.BASEPATH, constants.ASSETDIR_NAMESPACE
       )
+    full_path = os.path.join(base_path, path) if base_path else path
 
+    # Look in the dict of pre-loaded assets before checking the filesystem.
     if path in self._parent.namescope.assets:
-      # Look in the dict of pre-loaded assets before checking the filesystem.
       contents = self._parent.namescope.assets[path]
+    elif full_path in self._parent.namescope.assets:
+      contents = self._parent.namescope.assets[full_path]
     else:
-      # Construct the full path to the asset file, prefixed by the path to the
-      # model directory, and by `meshdir` or `texturedir` if appropriate.
-      path_parts = []
       if self._parent.namescope.model_dir:
-        path_parts.append(self._parent.namescope.model_dir)
-
-      if self._parent.namescope.has_identifier(
-          constants.BASEPATH, self._path_namespace
-      ):
-        base_path = self._parent.namescope.get(
-            constants.BASEPATH, self._path_namespace
-        )
-        path_parts.append(base_path)
-      elif (
-          self._path_namespace
-          in (constants.TEXTUREDIR_NAMESPACE, constants.MESHDIR_NAMESPACE)
-          and assetdir is not None
-      ):
-        path_parts.append(assetdir)
-      path_parts.append(path)
-      full_path = os.path.join(*path_parts)  # pylint: disable=no-value-for-parameter
+        full_path = os.path.join(self._parent.namescope.model_dir, full_path)
       contents = resources.GetResource(full_path)
 
     if self._parent.tag == constants.SKIN:

dm_control/mjcf/attribute_test.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ============================================================================
 
-"""Tests for `dm_control.mjcf.attribute`."""
-
 import contextlib
 import hashlib
 import os
@@ -455,6 +453,19 @@ def testFileFromAssetsDict(self):
         contents=contents, extension=extension, prefix=prefix)
     self.assertEqual(text_file.file, expected_value)
 
+  def testFileFromAssetsDictWithBasepath(self):
+    prefix = 'fake_filename'
+    extension = '.whatever'
+    path = 'invalid/path/' + prefix + extension
+    contents = 'Fake contents'
+    assets = {os.path.join(ASSETS_DIR, path): contents}
+    mujoco = element.RootElement(assets=assets)
+    mujoco.files.text_path = ASSETS_DIR
+    text_file = mujoco.files.add('text', file=path)
+    expected_value = attribute.Asset(
+        contents=contents, extension=extension, prefix=prefix)
+    self.assertEqual(text_file.file, expected_value)
+
   def testFileExceptions(self):
     mujoco = self._mujoco
     text_file = mujoco.files.add('text')