From 2b4975811b4273415675ad6ce9753adb1637996a Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Wed, 17 Apr 2024 05:57:57 -0700
Subject: [PATCH] Fix bug in Python binding of mj_saveModel

PiperOrigin-RevId: 625655183
Change-Id: I80d3390d739a8d6e280343087b79c15b543ff50a
---
 doc/changelog.rst          |  4 +++-
 python/mujoco/functions.cc | 16 ++++++++++------
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index d1c7c1ec..fb59d33f 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -19,7 +19,9 @@ MJX
 Bug fixes
 ^^^^^^^^^
 4. Defaults of lights were not being saved, now fixed.
-5. Prevent overwriting of frame names by body names when saving an XML.  Bug introduced in 3.1.4.
+5. Prevent overwriting of frame names by body names when saving an XML.  Introduced in 3.1.4.
+6. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
+   not optional.
 
 
 Version 3.1.4 (April 10th, 2024)
diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc
index 0282f8ca..8cb6b65c 100644
--- a/python/mujoco/functions.cc
+++ b/python/mujoco/functions.cc
@@ -106,11 +106,11 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mj_defaultOption>(pymodule);
   Def<traits::mj_defaultVisual>(pymodule);
   // Skipped: mj_copyModel (have MjModel.__copy__, memory managed by MjModel)
-  DEF_WITH_OMITTED_PY_ARGS(traits::mj_saveModel, "buffer_sz")(
-      pymodule,
-      [](const raw::MjModel* m, const std::optional<std::string>& filename,
+  pymodule.def(
+      "mj_saveModel",
+      [](const MjModelWrapper& m, const std::optional<std::string>& filename = std::nullopt,
          std::optional<
-             Eigen::Ref<Eigen::Vector<std::uint8_t, Eigen::Dynamic>>> buffer) {
+             Eigen::Ref<Eigen::Vector<std::uint8_t, Eigen::Dynamic>>> buffer = std::nullopt) {
         void* buffer_ptr = nullptr;
         int buffer_sz = 0;
         if (buffer.has_value()) {
@@ -118,9 +118,13 @@ PYBIND11_MODULE(_functions, pymodule) {
           buffer_sz = buffer->size();
         }
         return InterceptMjErrors(::mj_saveModel)(
-            m, filename.has_value() ? filename->c_str() : nullptr,
+            m.get(), filename.has_value() ? filename->c_str() : nullptr,
             buffer_ptr, buffer_sz);
-      });
+      },
+      py::arg("m"), py::arg_v("filename", std::nullopt),
+      py::arg_v("buffer", std::nullopt),
+      py::doc(traits::mj_saveModel::doc),
+      py::call_guard<py::gil_scoped_release>());
   // Skipped: mj_loadModel (have MjModel.from_binary_path)
   // Skipped: mj_deleteModel (have MjModel.__del__)
   Def<traits::mj_sizeModel>(pymodule);