Fixes and enhancements to _elementtree:

eliben · eliben · commit ebf37a2ffbe0 · 2012-04-03T22:02:37.000+03:00
* Fixed refleak problems when GC collection is run (see messages in issue #14065) * Added weakref support to Element objects
diff --git a/Lib/test/test_xml_etree.py b/Lib/test/test_xml_etree.py
@@ -1859,6 +1859,41 @@ class Dummy:
         gc_collect()
         self.assertIsNone(wref())
 
+        # A longer cycle: d->e->e2->d
+        e = ET.Element('joe')
+        d = Dummy()
+        d.dummyref = e
+        wref = weakref.ref(d)
+        e2 = ET.SubElement(e, 'foo', attr=d)
+        del d, e, e2
+        gc_collect()
+        self.assertIsNone(wref())
+
+        # A cycle between Element objects as children of one another
+        # e1->e2->e3->e1
+        e1 = ET.Element('e1')
+        e2 = ET.Element('e2')
+        e3 = ET.Element('e3')
+        e1.append(e2)
+        e2.append(e2)
+        e3.append(e1)
+        wref = weakref.ref(e1)
+        del e1, e2, e3
+        gc_collect()
+        self.assertIsNone(wref())
+
+    def test_weakref(self):
+        flag = False
+        def wref_cb(w):
+            nonlocal flag
+            flag = True
+        e = ET.Element('e')
+        wref = weakref.ref(e, wref_cb)
+        self.assertEqual(wref().tag, 'e')
+        del e
+        self.assertEqual(flag, True)
+        self.assertEqual(wref(), None)
+
 
 class ElementTreeTest(unittest.TestCase):
     def test_istype(self):
diff --git a/Modules/_elementtree.c b/Modules/_elementtree.c
@@ -48,6 +48,7 @@
 /* See http://www.python.org/psf/license for licensing details. */
 
 #include "Python.h"
+#include "structmember.h"
 
 #define VERSION "1.0.6"
 
@@ -229,6 +230,8 @@ typedef struct {
 
     ElementObjectExtra* extra;
 
+    PyObject *weakreflist; /* For tp_weaklistoffset */
+
 } ElementObject;
 
 static PyTypeObject Element_Type;
@@ -261,17 +264,24 @@ create_extra(ElementObject* self, PyObject* attrib)
 LOCAL(void)
 dealloc_extra(ElementObject* self)
 {
-    int i;
+    if (!self->extra)
+        return;
+
+    /* Avoid DECREFs calling into this code again (cycles, etc.)
+    */
+    ElementObjectExtra *myextra = self->extra;
+    self->extra = NULL;
 
-    Py_DECREF(self->extra->attrib);
+    Py_DECREF(myextra->attrib);
 
-    for (i = 0; i < self->extra->length; i++)
-        Py_DECREF(self->extra->children[i]);
+    int i;
+    for (i = 0; i < myextra->length; i++)
+        Py_DECREF(myextra->children[i]);
 
-    if (self->extra->children != self->extra->_children)
-        PyObject_Free(self->extra->children);
+    if (myextra->children != myextra->_children)
+        PyObject_Free(myextra->children);
 
-    PyObject_Free(self->extra);
+    PyObject_Free(myextra);
 }
 
 /* Convenience internal function to create new Element objects with the given
@@ -308,6 +318,8 @@ create_new_element(PyObject* tag, PyObject* attrib)
     Py_INCREF(Py_None);
     self->tail = Py_None;
 
+    self->weakreflist = NULL;
+
     ALLOC(sizeof(ElementObject), "create element");
     PyObject_GC_Track(self);
     return (PyObject*) self;
@@ -328,6 +340,7 @@ element_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         e->tail = Py_None;
 
         e->extra = NULL;
+        e->weakreflist = NULL;
     }
     return (PyObject *)e;
 }
@@ -576,26 +589,39 @@ element_gc_traverse(ElementObject *self, visitproc visit, void *arg)
 static int
 element_gc_clear(ElementObject *self)
 {
-    PyObject *text = JOIN_OBJ(self->text);
-    PyObject *tail = JOIN_OBJ(self->tail);
     Py_CLEAR(self->tag);
-    Py_CLEAR(text);
-    Py_CLEAR(tail);
 
-    /* After dropping all references from extra, it's no longer valid anyway,
-    ** so fully deallocate it (see also element_clearmethod)
+    /* The following is like Py_CLEAR for self->text and self->tail, but
+     * written explicitily because the real pointers hide behind access
+     * macros.
     */
-    if (self->extra) {
-        dealloc_extra(self);
-        self->extra = NULL;
+    if (self->text) {
+        PyObject *tmp = JOIN_OBJ(self->text);
+        self->text = NULL;
+        Py_DECREF(tmp);
+    }
+
+    if (self->tail) {
+        PyObject *tmp = JOIN_OBJ(self->tail);
+        self->tail = NULL;
+        Py_DECREF(tmp);
     }
+
+    /* After dropping all references from extra, it's no longer valid anyway,
+     * so fully deallocate it.
+    */
+    dealloc_extra(self);
     return 0;
 }
 
 static void
 element_dealloc(ElementObject* self)
 {
     PyObject_GC_UnTrack(self);
+
+    if (self->weakreflist != NULL)
+        PyObject_ClearWeakRefs((PyObject *) self);
+
     /* element_gc_clear clears all references and deallocates extra
     */
     element_gc_clear(self);
@@ -626,10 +652,7 @@ element_clearmethod(ElementObject* self, PyObject* args)
     if (!PyArg_ParseTuple(args, ":clear"))
         return NULL;
 
-    if (self->extra) {
-        dealloc_extra(self);
-        self->extra = NULL;
-    }
+    dealloc_extra(self);
 
     Py_INCREF(Py_None);
     Py_DECREF(JOIN_OBJ(self->text));
@@ -1693,7 +1716,7 @@ static PyTypeObject Element_Type = {
     (traverseproc)element_gc_traverse,              /* tp_traverse */
     (inquiry)element_gc_clear,                      /* tp_clear */
     0,                                              /* tp_richcompare */
-    0,                                              /* tp_weaklistoffset */
+    offsetof(ElementObject, weakreflist),           /* tp_weaklistoffset */
     0,                                              /* tp_iter */
     0,                                              /* tp_iternext */
     element_methods,                                /* tp_methods */
