Fix a memory leak when using Python callbacks.

Fixes #527.

PiperOrigin-RevId: 481173711
Change-Id: Ie176a8bda09727b6be8b6641d3be9a75c675d566

python/mujoco/callbacks.cc

@@ -46,7 +46,7 @@ using enable_if_not_const_t =
 // table that associates raw MuJoCo struct pointers back to the pointers to
 // their corresponding wrappers.
 template <typename Raw>
-static enable_if_not_const_t<Raw, py::handle> MjWrapperLookup(Raw* ptr) {
+static enable_if_not_const_t<Raw, py::object> MjWrapperLookup(Raw* ptr) {
   using LookupFnType = MjWrapper<Raw>* (Raw*);
   static LookupFnType* const lookup = []() -> LookupFnType* {
     py::gil_scoped_acquire gil;
@@ -93,8 +93,8 @@ static enable_if_not_const_t<Raw, py::handle> MjWrapperLookup(Raw* ptr) {
     const auto [src, type] =
         py::detail::type_caster_base<MjWrapper<Raw>>::src_and_type(wrapper);
     if (type) {
-      py::handle instance =
-          py::detail::find_registered_python_instance(wrapper, type);
+      py::object instance = py::reinterpret_steal<py::object>(
+          py::detail::find_registered_python_instance(wrapper, type));
       if (!instance) {
         if (!PyErr_Occurred()) {
           PyErr_SetString(
@@ -117,7 +117,7 @@ static enable_if_not_const_t<Raw, py::handle> MjWrapperLookup(Raw* ptr) {
 }
 
 template <typename Raw>
-static const py::handle MjWrapperLookup(const Raw* ptr) {
+static const py::object MjWrapperLookup(const Raw* ptr) {
   return MjWrapperLookup(const_cast<Raw*>(ptr));
 }
 
@@ -172,28 +172,28 @@ static PyObject* py_mju_user_free = nullptr;
 static PyObject* py_mjcb_passive = nullptr;
 static void PyMjcbPassive(const raw::MjModel* m, raw::MjData* d) {
   CallPyCallback<void>("mjcb_passive", py_mjcb_passive,
-                        MjWrapperLookup(m), MjWrapperLookup(d));
+                       MjWrapperLookup(m), MjWrapperLookup(d));
 }
 
 static PyObject* py_mjcb_control = nullptr;
 static void PyMjcbControl(const raw::MjModel* m, raw::MjData* d) {
   CallPyCallback<void>("mjcb_control", py_mjcb_control,
-                         MjWrapperLookup(m), MjWrapperLookup(d));
+                       MjWrapperLookup(m), MjWrapperLookup(d));
 }
 
 static PyObject* py_mjcb_contactfilter = nullptr;
 static int PyMjcbContactfilter(
     const raw::MjModel* m, raw::MjData* d, int geom1, int geom2) {
   return CallPyCallback<int>("mjcb_contactfilter", py_mjcb_contactfilter,
-                               MjWrapperLookup(m), MjWrapperLookup(d),
-                               geom1, geom2);
+                             MjWrapperLookup(m), MjWrapperLookup(d),
+                             geom1, geom2);
 }
 
 static PyObject* py_mjcb_sensor = nullptr;
 static void
 PyMjcbSensor(const raw::MjModel* m, raw::MjData* d, int stage) {
   CallPyCallback<void>("mjcb_sensor", py_mjcb_sensor,
-                         MjWrapperLookup(m), MjWrapperLookup(d), stage);
+                       MjWrapperLookup(m), MjWrapperLookup(d), stage);
 }
 
 static PyObject* py_mjcb_time = nullptr;