From c6984b42af6b032c5c9ec8e548c7c02ce3680e46 Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Fri, 18 Mar 2022 13:00:32 -0700
Subject: [PATCH] Add a free() method to mujoco.MjrContext.

Up the mujoco bindings version to 2.1.2.post2.

Keep the capsule deconstructor which calls mj_freeContext, because mjr_freeContext is idempotent.

PiperOrigin-RevId: 435714183
Change-Id: I9f8615f85f2b3e383f447f6684e61200039d57d4
---
 python/mujoco/render.cc      | 12 ++++++++++++
 python/mujoco/render_test.py | 25 +++++++++++++++++++------
 python/setup.py              |  2 +-
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/python/mujoco/render.cc b/python/mujoco/render.cc
index fd7b9ffd..22699f24 100644
--- a/python/mujoco/render.cc
+++ b/python/mujoco/render.cc
@@ -37,6 +37,8 @@ class MjWrapper<raw::MjrContext> : public WrapperBase<raw::MjrContext> {
   MjWrapper(MjWrapper&&) = default;
   ~MjWrapper() = default;
 
+  void Free();
+
 #define X(var) py_array_or_tuple_t<mjtNum> var
   X(fogRGBA);
   X(auxWidth);
@@ -121,6 +123,12 @@ MjrContextWrapper::MjWrapper(const MjModelWrapper& model, int fontscale)
       X(charWidthBig) {}
 #undef X_SKIN
 #undef X
+
+void MjrContextWrapper::Free() {
+  // mjr_freeContext is safe to call multiple times.
+  InterceptMjErrors(mjr_freeContext)(ptr_);
+}
+
 }  // namespace _impl
 
 namespace {
@@ -156,6 +164,10 @@ PYBIND11_MODULE(_render, pymodule) {
   py::class_<MjrContextWrapper> mjrContext(pymodule, "MjrContext");
   mjrContext.def(py::init<>());
   mjrContext.def(py::init<const MjModelWrapper&, int>());
+  mjrContext.def(
+      "free", [](MjrContextWrapper& self) { self.Free(); },
+      py::doc("Frees resources in current active OpenGL context, sets struct "
+              "to default."));
 #define X(var)                                                       \
   mjrContext.def_property(                                           \
       #var, [](const MjrContextWrapper& c) { return c.get()->var; }, \
diff --git a/python/mujoco/render_test.py b/python/mujoco/render_test.py
index 4c2f1b61..5f27a7a8 100644
--- a/python/mujoco/render_test.py
+++ b/python/mujoco/render_test.py
@@ -49,13 +49,10 @@ class MuJoCoRenderTest(absltest.TestCase):
     scene = mujoco.MjvScene(self.model, maxgeom=0)
     mujoco.mjv_updateScene(
         self.model, self.data, mujoco.MjvOption(), mujoco.MjvPerturb(),
-        mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL.value, scene)
+        mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
 
-    context = mujoco.MjrContext(
-        self.model,
-        mujoco.mjtFontScale.mjFONTSCALE_150.value)
-    mujoco.mjr_setBuffer(
-        mujoco.mjtFramebuffer.mjFB_OFFSCREEN.value, context)
+    context = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
+    mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, context)
 
     # MuJoCo's default render buffer size is 640x480.
     full_rect = mujoco.MjrRect(0, 0, 640, 480)
@@ -76,6 +73,22 @@ class MuJoCoRenderTest(absltest.TestCase):
     mujoco.mjr_readPixels(
         np.reshape(upside_down_image, -1), None, full_rect, context)
     np.testing.assert_array_equal(upside_down_image, expected_upside_down_image)
+    context.free()
+
+  def test_safe_to_free_context_twice(self):
+    self.model = mujoco.MjModel.from_xml_string('<mujoco><worldbody/></mujoco>')
+    self.data = mujoco.MjData(self.model)
+
+    scene = mujoco.MjvScene(self.model, maxgeom=0)
+    mujoco.mjv_updateScene(
+        self.model, self.data, mujoco.MjvOption(), None,
+        mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
+
+    context = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
+    mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, context)
+
+    context.free()
+    context.free()
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/python/setup.py b/python/setup.py
index 4cab1805..3350af46 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -30,7 +30,7 @@ from setuptools import find_packages
 from setuptools import setup
 from setuptools.command import build_ext
 
-__version__ = '2.1.2.post1'
+__version__ = '2.1.2.post2'
 
 MUJOCO_CMAKE = 'MUJOCO_CMAKE'
 MUJOCO_CMAKE_ARGS = 'MUJOCO_CMAKE_ARGS'