Issue #14065: Added cyclic GC support to ET.Element

eliben · eliben · commit 0192ba33b4ff · 2012-03-30T16:38:33.000+03:00
diff --git a/Lib/test/test_xml_etree.py b/Lib/test/test_xml_etree.py
@@ -14,9 +14,10 @@
 # Don't re-import "xml.etree.ElementTree" module in the docstring,
 # except if the test is specific to the Python implementation.
 
-import sys
+import gc
 import html
 import io
+import sys
 import unittest
 
 from test import support
@@ -1846,6 +1847,30 @@ def test_augmentation_type_errors(self):
         self.assertRaises(TypeError, e.extend, [ET.Element('bar'), 'foo'])
         self.assertRaises(TypeError, e.insert, 0, 'foo')
 
+    def test_cyclic_gc(self):
+        class ShowGC:
+            def __init__(self, flaglist):
+                self.flaglist = flaglist
+            def __del__(self):
+                self.flaglist.append(1)
+
+        # Test the shortest cycle: lst->element->lst
+        fl = []
+        lst = [ShowGC(fl)]
+        lst.append(ET.Element('joe', attr=lst))
+        del lst
+        gc.collect()
+        self.assertEqual(fl, [1])
+
+        # A longer cycle: lst->e->e2->lst
+        fl = []
+        e = ET.Element('joe')
+        lst = [ShowGC(fl), e]
+        e2 = ET.SubElement(e, 'foo', attr=lst)
+        del lst, e, e2
+        gc.collect()
+        self.assertEqual(fl, [1])
+
 
 class ElementTreeTest(unittest.TestCase):
     def test_istype(self):
diff --git a/Modules/_elementtree.c b/Modules/_elementtree.c
@@ -282,7 +282,7 @@ create_new_element(PyObject* tag, PyObject* attrib)
 {
     ElementObject* self;
 
-    self = PyObject_New(ElementObject, &Element_Type);
+    self = PyObject_GC_New(ElementObject, &Element_Type);
     if (self == NULL)
         return NULL;
 
@@ -309,7 +309,7 @@ create_new_element(PyObject* tag, PyObject* attrib)
     self->tail = Py_None;
 
     ALLOC(sizeof(ElementObject), "create element");
-
+    PyObject_GC_Track(self);
     return (PyObject*) self;
 }
 
@@ -556,19 +556,51 @@ subelement(PyObject* self, PyObject* args, PyObject* kw)
     return elem;
 }
 
-static void
-element_dealloc(ElementObject* self)
+static int
+element_gc_traverse(ElementObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(self->tag);
+    Py_VISIT(JOIN_OBJ(self->text));
+    Py_VISIT(JOIN_OBJ(self->tail));
+
+    if (self->extra) {
+        int i;
+        Py_VISIT(self->extra->attrib);
+
+        for (i = 0; i < self->extra->length; ++i)
+            Py_VISIT(self->extra->children[i]);
+    }
+    return 0;
+}
+
+static int
+element_gc_clear(ElementObject *self)
 {
-    if (self->extra)
+    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)
+    */
+    if (self->extra) {
         dealloc_extra(self);
+        self->extra = NULL;
+    }
+    return 0;
+}
 
-    /* discard attributes */
-    Py_DECREF(self->tag);
-    Py_DECREF(JOIN_OBJ(self->text));
-    Py_DECREF(JOIN_OBJ(self->tail));
+static void
+element_dealloc(ElementObject* self)
+{
+    PyObject_GC_UnTrack(self);
+    /* element_gc_clear clears all references and deallocates extra
+    */
+    element_gc_clear(self);
 
     RELEASE(sizeof(ElementObject), "destroy element");
-
     Py_TYPE(self)->tp_free((PyObject *)self);
 }
 
@@ -589,7 +621,7 @@ element_append(ElementObject* self, PyObject* args)
 }
 
 static PyObject*
-element_clear(ElementObject* self, PyObject* args)
+element_clearmethod(ElementObject* self, PyObject* args)
 {
     if (!PyArg_ParseTuple(args, ":clear"))
         return NULL;
@@ -1505,7 +1537,7 @@ element_ass_subscr(PyObject* self_, PyObject* item, PyObject* value)
 
 static PyMethodDef element_methods[] = {
 
-    {"clear", (PyCFunction) element_clear, METH_VARARGS},
+    {"clear", (PyCFunction) element_clearmethod, METH_VARARGS},
 
     {"get", (PyCFunction) element_get, METH_VARARGS},
     {"set", (PyCFunction) element_set, METH_VARARGS},
@@ -1655,10 +1687,11 @@ static PyTypeObject Element_Type = {
     (getattrofunc)element_getattro,                 /* tp_getattro */
     0,                                              /* tp_setattro */
     0,                                              /* tp_as_buffer */
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,       /* tp_flags */
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
+                                                    /* tp_flags */
     0,                                              /* tp_doc */
-    0,                                              /* tp_traverse */
-    0,                                              /* tp_clear */
+    (traverseproc)element_gc_traverse,              /* tp_traverse */
+    (inquiry)element_gc_clear,                      /* tp_clear */
     0,                                              /* tp_richcompare */
     0,                                              /* tp_weaklistoffset */
     0,                                              /* tp_iter */
