Update spec.from_string() and spec.from_file() to match the respective mjModel functions.

PiperOrigin-RevId: 678230403
Change-Id: I7f3d8c5fd44d30b5695eee1f8513cc27f4158d94

python/mujoco/specs.cc

@@ -19,6 +19,7 @@
 #include <optional>
 #include <string>
 #include <string_view>  // IWYU pragma: keep
+#include <unordered_map>
 #include <vector>       // IWYU pragma: keep
 
 #include <Eigen/Core>
@@ -73,6 +74,37 @@ struct MjSpec {
   raw::MjSpec* ptr;
 };
 
+template <typename LoadFunc>
+static raw::MjSpec* LoadSpecFileImpl(
+    const std::string& filename,
+    const std::vector<_impl::VfsAsset>& assets,
+    LoadFunc&& loadfunc) {
+  mjVFS vfs;
+  mjVFS* vfs_ptr = nullptr;
+  if (!assets.empty()) {
+    mj_defaultVFS(&vfs);
+    vfs_ptr = &vfs;
+    for (const auto& asset : assets) {
+      std::string buffer_name = _impl::StripPath(asset.name);
+      const int vfs_error = InterceptMjErrors(mj_addBufferVFS)(
+          vfs_ptr, buffer_name.c_str(), asset.content, asset.content_size);
+      if (vfs_error) {
+        mj_deleteVFS(vfs_ptr);
+        if (vfs_error == 2) {
+          throw py::value_error("Repeated file name in assets dict: " +
+                                buffer_name);
+        } else {
+          throw py::value_error("Asset failed to load: " + buffer_name);
+        }
+      }
+    }
+  }
+
+  raw::MjSpec* spec = loadfunc(filename.c_str(), vfs_ptr);
+  mj_deleteVFS(vfs_ptr);
+  return spec;
+}
+
 PYBIND11_MODULE(_specs, m) {
   auto structs_m = py::module::import("mujoco._structs");
   py::function mjmodel_from_spec_ptr =
@@ -196,27 +228,25 @@ PYBIND11_MODULE(_specs, m) {
   mjSpec.def(
       "from_file",
       [](MjSpec& self, std::string& filename,
-         std::optional<py::dict>& assets) -> void {
-        mjVFS vfs;
-        mj_defaultVFS(&vfs);
-        if (assets.has_value()) {
-          for (auto item : assets.value()) {
-            std::string buffer = py::cast<std::string>(item.second);
-            mj_addBufferVFS(&vfs, py::cast<std::string>(item.first).c_str(),
-                            buffer.c_str(), buffer.size());
-          };
-        }
-        std::array<char, 1024> err;
-        err[0] = '\0';
-        mj_deleteSpec(self.ptr);
-        self.ptr = mj_parseXML(filename.c_str(), &vfs, err.data(), err.size());
-        mj_deleteVFS(&vfs);
-        if (!self.ptr) {
-          throw FatalError(std::string(err.data()));
+         std::optional<std::unordered_map<std::string, py::bytes>>& assets)
+          -> void {
+        const auto converted_assets = _impl::ConvertAssetsDict(assets);
+        {
+          py::gil_scoped_release no_gil;
+          char error[1024];
+          mj_deleteSpec(self.ptr);
+          self.ptr = LoadSpecFileImpl(
+              filename, converted_assets,
+              [&error](const char* filename, const mjVFS* vfs) {
+                return InterceptMjErrors(mj_parseXML)(
+                    filename, vfs, error, sizeof(error));
+              });
+          if (!self.ptr) {
+            throw py::value_error(error);
+          }
         }
       },
-      py::arg("filename"),
-      py::arg("assets") = py::none(), R"mydelimiter(
+      py::arg("filename"), py::arg("assets") = py::none(), R"mydelimiter(
     Creates a spec from an XML file.
 
     Parameters
@@ -230,27 +260,34 @@ PYBIND11_MODULE(_specs, m) {
   mjSpec.def(
       "from_string",
       [](MjSpec& self, std::string& xml,
-         std::optional<py::dict>& assets) -> void {
-        mjVFS vfs;
-        mj_defaultVFS(&vfs);
-        if (assets.has_value()) {
-          for (auto item : assets.value()) {
-            std::string buffer = py::cast<std::string>(item.second);
-            mj_addBufferVFS(&vfs, py::cast<std::string>(item.first).c_str(),
-                            buffer.c_str(), buffer.size());
-          };
-        }
-        std::array<char, 1024> err;
-        err[0] = '\0';
-        mj_deleteSpec(self.ptr);
-        self.ptr = mj_parseXMLString(xml.c_str(), &vfs, err.data(), err.size());
-        mj_deleteVFS(&vfs);
-        if (!self.ptr) {
-          throw FatalError(std::string(err.data()));
+         std::optional<std::unordered_map<std::string, py::bytes>>& assets)
+          -> void {
+        auto converted_assets = _impl::ConvertAssetsDict(assets);
+        {
+          py::gil_scoped_release no_gil;
+          std::string model_filename = "model_.xml";
+          if (assets.has_value()) {
+            while (assets->find(model_filename) != assets->end()) {
+              model_filename =
+                  model_filename.substr(0, model_filename.size() - 4) + "_.xml";
+            }
+          }
+          converted_assets.emplace_back(
+              model_filename.c_str(), xml.c_str(), xml.length());
+          char error[1024];
+          mj_deleteSpec(self.ptr);
+          self.ptr = LoadSpecFileImpl(
+              model_filename, converted_assets,
+              [&error](const char* filename, const mjVFS* vfs) {
+                return InterceptMjErrors(mj_parseXML)(
+                    filename, vfs, error, sizeof(error));
+              });
+          if (!self.ptr) {
+            throw py::value_error(error);
+          }
         }
       },
-      py::arg("xml"),
-      py::arg("assets") = py::none(), R"mydelimiter(
+      py::arg("xml"), py::arg("assets") = py::none(), R"mydelimiter(
     Creates a spec from an XML string.
 
     Parameters

python/mujoco/specs_test.py

@@ -611,6 +611,25 @@ class SpecsTest(absltest.TestCase):
     model = spec.compile({'cube.obj': cube})
     self.assertEqual(model.nmeshvert, 8)
 
+  def test_include(self):
+    included_xml = """
+      <mujoco>
+        <worldbody>
+          <body>
+            <geom type="box" size="1 1 1"/>
+          </body>
+        </worldbody>
+      </mujoco>
+    """
+    spec = mujoco.MjSpec()
+    spec.from_string(textwrap.dedent("""
+      <mujoco model="MuJoCo Model">
+        <include file="included.xml"/>
+      </mujoco>
+    """), {'included.xml': included_xml.encode('utf-8')})
+    self.assertEqual(spec.worldbody.first_body().first_geom().type,
+                     mujoco.mjtGeom.mjGEOM_BOX)
+
   def test_delete(self):
     filename = '../../test/testdata/model.xml'
 

python/mujoco/structs.cc

@@ -87,21 +87,6 @@ inline std::size_t NConMax(const mjData* d) {
   return d->narena / sizeof(mjContact);
 }
 
-// strip path prefix from filename and make lowercase
-std::string StripPath(const char* name) {
-  std::string filename(name);
-  size_t start = filename.find_last_of("/\\");
-
-  // get name without path
-  if (start != std::string::npos) {
-    filename = filename.substr(start + 1, filename.size() - start - 1);
-  }
-
-  // make lowercase
-  std::transform(filename.begin(), filename.end(), filename.begin(),
-                 [](unsigned char c) { return std::tolower(c); });
-  return filename;
-}
 }  // namespace
 
 // ==================== MJOPTION ===============================================
@@ -318,16 +303,6 @@ MjModelWrapper::~MjWrapper() {
   }
 }
 
-namespace {
-struct VfsAsset {
-  VfsAsset(const char* name, const void* content, std::size_t content_size)
-      : name(name), content(content), content_size(content_size) {}
-  const char* name;
-  const void* content;
-  std::size_t content_size;
-};
-}
-
 // Helper function for both LoadXMLFile and LoadBinaryFile.
 // Creates a temporary MJB from the assets dictionary if one is supplied.
 template <typename LoadFunc>
@@ -365,22 +340,6 @@ static raw::MjModel* LoadModelFileImpl(
   return model;
 }
 
-// Converts a dict with py::bytes value to a vector of standard C++ types.
-// This allows us to release the GIL early. Note that the vector consists only
-// of pointers to existing data so no substantial data copies are being made.
-static std::vector<VfsAsset>
-ConvertAssetsDict(
-    const std::optional<std::unordered_map<std::string, py::bytes>>& assets) {
-  std::vector<VfsAsset> out;
-  if (assets.has_value()) {
-    for (const auto& [name, content] : *assets) {
-      out.emplace_back(name.c_str(), PYBIND11_BYTES_AS_STRING(content.ptr()),
-                       py::len(content));
-    }
-  }
-  return out;
-}
-
 MjModelWrapper MjModelWrapper::LoadXMLFile(
     const std::string& filename,
     const std::optional<std::unordered_map<std::string, py::bytes>>& assets) {

python/mujoco/structs.h

@@ -15,7 +15,9 @@
 #ifndef MUJOCO_PYTHON_STRUCTS_H_
 #define MUJOCO_PYTHON_STRUCTS_H_
 
+#include <algorithm>
 #include <array>
+#include <cctype>
 #include <cstddef>
 #include <functional>
 #include <istream>
@@ -39,9 +41,51 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/pytypes.h>
 
+namespace py = ::pybind11;
+
 namespace mujoco::python {
 namespace _impl {
 
+struct VfsAsset {
+  VfsAsset(const char* name, const void* content, std::size_t content_size)
+      : name(name), content(content), content_size(content_size) {}
+  const char* name;
+  const void* content;
+  std::size_t content_size;
+};
+
+// strip path prefix from filename and make lowercase
+inline std::string StripPath(const char* name) {
+  std::string filename(name);
+  size_t start = filename.find_last_of("/\\");
+
+  // get name without path
+  if (start != std::string::npos) {
+    filename = filename.substr(start + 1, filename.size() - start - 1);
+  }
+
+  // make lowercase
+  std::transform(filename.begin(), filename.end(), filename.begin(),
+                 [](unsigned char c) { return std::tolower(c); });
+  return filename;
+}
+
+
+// Converts a dict with py::bytes value to a vector of standard C++ types.
+// This allows us to release the GIL early. Note that the vector consists only
+// of pointers to existing data so no substantial data copies are being made.
+inline std::vector<VfsAsset> ConvertAssetsDict(
+    const std::optional<std::unordered_map<std::string, py::bytes>>& assets) {
+  std::vector<VfsAsset> out;
+  if (assets.has_value()) {
+    for (const auto& [name, content] : *assets) {
+      out.emplace_back(name.c_str(), PYBIND11_BYTES_AS_STRING(content.ptr()),
+                       py::len(content));
+    }
+  }
+  return out;
+}
+
 template <typename T>
 class WrapperBase {
  public: