Bug fix for:

[ #5000 ] Changin EditMode in Script wrecks memory

Armatures create weakreferences in __main__.__dict__.  When Window.Editmode is called, the weaklist it iterated over and armatures are updated.

3 files changed, 51 insertions, 5 deletions

diff --git a/source/blender/python/api2_2x/Armature.c b/source/blender/python/api2_2x/Armature.c
index f542caeb872..37c28b532d9 100644
--- a/source/blender/python/api2_2x/Armature.c
+++ b/source/blender/python/api2_2x/Armature.c
@@ -1111,6 +1111,8 @@ static PyObject *Armature_repr(BPy_Armature *self)
 ///tp_dealloc
 static void Armature_dealloc(BPy_Armature * self)
 {
+	if (self->weaklist != NULL)
+        PyObject_ClearWeakRefs((PyObject *) self);
 	Py_DECREF(self->Bones);
 	Armature_Type.tp_free(self);
 	return;
@@ -1137,12 +1139,12 @@ PyTypeObject Armature_Type = {
 	0,								//tp_getattro
 	0,								//tp_setattro
 	0,								//tp_as_buffer
-	Py_TPFLAGS_DEFAULT,				//tp_flags
+	Py_TPFLAGS_DEFAULT| Py_TPFLAGS_HAVE_WEAKREFS,				//tp_flags
 	BPy_Armature_doc,				//tp_doc
 	0,								//tp_traverse
 	0,								//tp_clear
 	0, 								//tp_richcompare
-	0,								//tp_weaklistoffset
+	offsetof(BPy_Armature, weaklist),	//tp_weaklistoffset
 	0,								//tp_iter
 	0,								//tp_iternext
 	BPy_Armature_methods,			//tp_methods
@@ -1328,14 +1330,27 @@ struct PyMethodDef M_Armature_methods[] = {
 	{NULL, NULL, 0, NULL}
 };
 //------------------VISIBLE PROTOTYPE IMPLEMENTATION-----------------------
+//------------------------Armature_RebuildEditbones (internal)
+PyObject * Armature_RebuildEditbones(PyObject *pyarmature)
+{
+	return Armature_makeEditable((BPy_Armature*)pyarmature);
+}
+
+//------------------------Armature_RebuildBones (internal)
+PyObject *Armature_RebuildBones(PyObject *pyarmature)
+{
+	return Armature_update((BPy_Armature*)pyarmature);
+}
 //-----------------(internal)
 //Converts a bArmature to a PyArmature
 PyObject *PyArmature_FromArmature(struct bArmature *armature)
 {
 	BPy_Armature *py_armature = NULL;
+	PyObject *maindict = NULL, *armdict = NULL, *weakref = NULL;
 
 	//create armature type
 	py_armature = (BPy_Armature*)Armature_Type.tp_alloc(&Armature_Type, 0); //*new*
+	py_armature->weaklist = NULL; //init the weaklist
 	if (!py_armature)
 		goto RuntimeError;
 	py_armature->armature = armature;
@@ -1345,6 +1360,14 @@ PyObject *PyArmature_FromArmature(struct bArmature *armature)
 	if (!py_armature->Bones)
 		goto RuntimeError;
 
+	//put a weakreference in __main__
+	maindict= PyModule_GetDict(PyImport_AddModule(	"__main__"));
+	armdict = PyDict_GetItemString(maindict, "armatures");
+	weakref = PyWeakref_NewProxy((PyObject*)py_armature, Py_None);
+	if (PyList_Append(armdict, weakref) == -1){
+		goto RuntimeError;
+	}
+
 	return (PyObject *) py_armature; 
 
 RuntimeError:
diff --git a/source/blender/python/api2_2x/Armature.h b/source/blender/python/api2_2x/Armature.h
index 8115207dc3b..f29254e9aa6 100644
--- a/source/blender/python/api2_2x/Armature.h
+++ b/source/blender/python/api2_2x/Armature.h
@@ -56,11 +56,14 @@ typedef struct {
 	PyObject_HEAD 
 	struct bArmature * armature;
 	BPy_BonesDict *Bones;          //BPy_BonesDict
+	PyObject *weaklist;
 } BPy_Armature;
 
 //-------------------VISIBLE PROTOTYPES-------------------------
 PyObject *PyArmature_FromArmature(struct bArmature *armature);
 struct bArmature *PyArmature_AsArmature(BPy_Armature *py_armature);
+PyObject * Armature_RebuildEditbones(PyObject *pyarmature);
+PyObject *Armature_RebuildBones(PyObject *pyarmature);
 
 int Armature_CheckPyObject( PyObject * pyobj );
 struct bArmature  *Armature_FromPyObject( PyObject * py_obj );
diff --git a/source/blender/python/api2_2x/Window.c b/source/blender/python/api2_2x/Window.c
index de9baec0a81..0bed7d9241d 100644
--- a/source/blender/python/api2_2x/Window.c
+++ b/source/blender/python/api2_2x/Window.c
@@ -57,7 +57,7 @@
 #include "Mathutils.h"
 #include "constant.h"
 #include "gen_utils.h"
-
+#include "Armature.h"
 
 /* See Draw.c */
 extern int EXPP_disable_force_draw;
@@ -890,6 +890,9 @@ static PyObject *M_Window_EditMode( PyObject * self, PyObject * args )
 	char *undo_str = "From script";
 	int undo_str_len = 11;
 	int do_undo = 1;
+	PyObject *maindict = NULL, *armlist = NULL;
+	PyObject *pyarmature = NULL;
+	int x;
 
 	if( !PyArg_ParseTuple( args,
 				"|hs#i", &status, &undo_str, &undo_str_len, &do_undo ) )
@@ -899,14 +902,31 @@ static PyObject *M_Window_EditMode( PyObject * self, PyObject * args )
 	if( status >= 0 ) {
 		if( status ) {
 			if( !G.obedit )
+				//update armatures
+				maindict= PyModule_GetDict(PyImport_AddModule(	"__main__"));
+				armlist = PyDict_GetItemString(maindict, "armatures");
+				for (x = 0; x < PySequence_Size(armlist); x++){
+					pyarmature = PyWeakref_GetObject(PySequence_GetItem(	armlist, x));
+					if (pyarmature != Py_None)
+						Armature_RebuildEditbones(pyarmature);
+				}
+				//enter editmode
 				enter_editmode(0);
 		} else if( G.obedit ) {
-			
 			if( undo_str_len > 63 )
 				undo_str[63] = '\0';	/* 64 is max */
-			
 			BIF_undo_push( undo_str ); /* This checks user undo settings */
 			exit_editmode( EM_FREEDATA );
+
+			//update armatures
+			maindict= PyModule_GetDict(PyImport_AddModule(	"__main__"));
+			armlist = PyDict_GetItemString(maindict, "armatures");
+			for (x = 0; x < PySequence_Size(armlist); x++){
+				pyarmature = PyWeakref_GetObject(PySequence_GetItem(	armlist, x));
+				if (pyarmature != Py_None)
+					Armature_RebuildBones(pyarmature);
+			}
+
 		}
 	}