PyAPI: draw handlers now store args in the capsule context

Was using the handlers custom-data which crashes
when Blender frees the screen data before Python removes the handler.

5 files changed, 17 insertions, 32 deletions

diff --git a/source/blender/editors/include/ED_space_api.h b/source/blender/editors/include/ED_space_api.h
index ddd8b59c264..709af6e8b09 100644
--- a/source/blender/editors/include/ED_space_api.h
+++ b/source/blender/editors/include/ED_space_api.h
@@ -75,7 +75,6 @@ void *ED_region_draw_cb_activate(struct ARegionType *,
                                  void *custumdata, int type);
 void ED_region_draw_cb_draw(const struct bContext *, struct ARegion *, int);
 void ED_region_draw_cb_exit(struct ARegionType *, void *);
-void *ED_region_draw_cb_customdata(void *handle);
 /* generic callbacks */
 /* ed_util.c */
 void ED_region_draw_mouse_line_cb(const struct bContext *C, struct ARegion *ar, void *arg_info);
diff --git a/source/blender/editors/space_api/spacetypes.c b/source/blender/editors/space_api/spacetypes.c
index 91dff98f906..3e343d5e35a 100644
--- a/source/blender/editors/space_api/spacetypes.c
+++ b/source/blender/editors/space_api/spacetypes.c
@@ -266,11 +266,6 @@ void ED_region_draw_cb_exit(ARegionType *art, void *handle)
 	}
 }
 
-void *ED_region_draw_cb_customdata(void *handle)
-{
-	return ((RegionDrawCB *)handle)->customdata;
-}
-
 void ED_region_draw_cb_draw(const bContext *C, ARegion *ar, int type)
 {
 	RegionDrawCB *rdc;
diff --git a/source/blender/python/intern/bpy_rna_callback.c b/source/blender/python/intern/bpy_rna_callback.c
index fd475ba503f..3a20ba385c6 100644
--- a/source/blender/python/intern/bpy_rna_callback.c
+++ b/source/blender/python/intern/bpy_rna_callback.c
@@ -296,7 +296,6 @@ PyObject *pyrna_callback_classmethod_add(PyObject *UNUSED(self), PyObject *args)
 		        wm,
 		        params.space_type, params.region_type,
 		        NULL, cb_wm_cursor_draw, (void *)args);
-		Py_INCREF(args);
 	}
 	else if (RNA_struct_is_a(srna, &RNA_Space)) {
 		const char *error_prefix = "Space.draw_handler_add";
@@ -343,7 +342,6 @@ PyObject *pyrna_callback_classmethod_add(PyObject *UNUSED(self), PyObject *args)
 					return NULL;
 				}
 				handle = ED_region_draw_cb_activate(art, cb_region_draw, (void *)args, params.event);
-				Py_INCREF(args);
 			}
 		}
 	}
@@ -352,7 +350,14 @@ PyObject *pyrna_callback_classmethod_add(PyObject *UNUSED(self), PyObject *args)
 		return NULL;
 	}
 
-	return PyCapsule_New((void *)handle, RNA_CAPSULE_ID, NULL);
+	PyObject *ret = PyCapsule_New((void *)handle, RNA_CAPSULE_ID, NULL);
+
+	/* Store 'args' in context as well as the handler custom-data,
+	 * because the handle may be freed by Blender (new file, new window... etc) */
+	PyCapsule_SetContext(ret, args);
+	Py_INCREF(args);
+
+	return ret;
 }
 
 PyObject *pyrna_callback_classmethod_remove(PyObject *UNUSED(self), PyObject *args)
@@ -360,8 +365,8 @@ PyObject *pyrna_callback_classmethod_remove(PyObject *UNUSED(self), PyObject *ar
 	PyObject *cls;
 	PyObject *py_handle;
 	void *handle;
-	void *customdata;
 	StructRNA *srna;
+	bool capsule_clear = false;
 
 	if (PyTuple_GET_SIZE(args) < 2) {
 		PyErr_SetString(PyExc_ValueError, "callback_remove(handler): expected at least 2 args");
@@ -378,6 +383,7 @@ PyObject *pyrna_callback_classmethod_remove(PyObject *UNUSED(self), PyObject *ar
 		PyErr_SetString(PyExc_ValueError, "callback_remove(handler): NULL handler given, invalid or already removed");
 		return NULL;
 	}
+	PyObject *handle_args = PyCapsule_GetContext(py_handle);
 
 	if (srna == &RNA_WindowManager) {
 		if (!PyArg_ParseTuple(
@@ -388,18 +394,8 @@ PyObject *pyrna_callback_classmethod_remove(PyObject *UNUSED(self), PyObject *ar
 		}
 		bContext *C = BPy_GetContext();
 		struct wmWindowManager *wm = CTX_wm_manager(C);
-
-		if (BLI_findindex(&wm->paintcursors, handle) == -1) {
-			/* FIXME(campbell): window manager has freed cursor, need to resolve refcount leak. */
-		}
-		else {
-			customdata = WM_paint_cursor_customdata_get(handle);
-			if (!WM_paint_cursor_end(wm, handle)) {
-				PyErr_SetString(PyExc_ValueError, "draw_cursor_remove(handler): cursor wasn't found");
-				return NULL;
-			}
-			Py_DECREF((PyObject *)customdata);
-		}
+		WM_paint_cursor_end(wm, handle);
+		capsule_clear = true;
 	}
 	else if (RNA_struct_is_a(srna, &RNA_Space)) {
 		const char *error_prefix = "Space.draw_handler_remove";
@@ -417,9 +413,6 @@ PyObject *pyrna_callback_classmethod_remove(PyObject *UNUSED(self), PyObject *ar
 			return NULL;
 		}
 
-		customdata = ED_region_draw_cb_customdata(handle);
-		Py_DECREF((PyObject *)customdata);
-
 		if (pyrna_enum_value_from_id(
 		            rna_enum_region_type_items, params.region_type_str,
 		            &params.region_type, error_prefix) == -1)
@@ -440,6 +433,7 @@ PyObject *pyrna_callback_classmethod_remove(PyObject *UNUSED(self), PyObject *ar
 					return NULL;
 				}
 				ED_region_draw_cb_exit(art, handle);
+				capsule_clear = true;
 			}
 		}
 	}
@@ -449,7 +443,10 @@ PyObject *pyrna_callback_classmethod_remove(PyObject *UNUSED(self), PyObject *ar
 	}
 
 	/* don't allow reuse */
-	PyCapsule_SetName(py_handle, RNA_CAPSULE_ID_INVALID);
+	if (capsule_clear) {
+		Py_DECREF(handle_args);
+		PyCapsule_SetName(py_handle, RNA_CAPSULE_ID_INVALID);
+	}
 
 	Py_RETURN_NONE;
 }
diff --git a/source/blender/windowmanager/WM_api.h b/source/blender/windowmanager/WM_api.h
index 4d3be1e3f3b..7628217d88f 100644
--- a/source/blender/windowmanager/WM_api.h
+++ b/source/blender/windowmanager/WM_api.h
@@ -179,7 +179,6 @@ struct wmPaintCursor *WM_paint_cursor_activate(
         void *customdata);
 
 bool		WM_paint_cursor_end(struct wmWindowManager *wm, struct wmPaintCursor *handle);
-void       *WM_paint_cursor_customdata_get(struct wmPaintCursor *pc);
 void		WM_paint_cursor_tag_redraw(struct wmWindow *win, struct ARegion *ar);
 
 
diff --git a/source/blender/windowmanager/intern/wm_operators.c b/source/blender/windowmanager/intern/wm_operators.c
index e12107283a5..3be213796d8 100644
--- a/source/blender/windowmanager/intern/wm_operators.c
+++ b/source/blender/windowmanager/intern/wm_operators.c
@@ -1878,11 +1878,6 @@ bool WM_paint_cursor_end(wmWindowManager *wm, wmPaintCursor *handle)
 	return false;
 }
 
-void *WM_paint_cursor_customdata_get(wmPaintCursor *pc)
-{
-	return pc->customdata;
-}
-
 /* *********************** radial control ****************** */
 
 #define WM_RADIAL_CONTROL_DISPLAY_SIZE (200 * UI_DPI_FAC)