Issue #14007: make TreeBuilder an actual type exposed from _elementtree, and subclassable.

Lib/test/test_xml_etree.py

@@ -1979,6 +1979,20 @@ class TreeBuilderTest(unittest.TestCase):
         parser.feed(self.sample1)
         self.assertIsNone(parser.close())
 
+    def test_subclass(self):
+        class MyTreeBuilder(ET.TreeBuilder):
+            def foobar(self, x):
+                return x * 2
+
+        tb = MyTreeBuilder()
+        self.assertEqual(tb.foobar(10), 20)
+
+        parser = ET.XMLParser(target=tb)
+        parser.feed(self.sample1)
+
+        e = parser.close()
+        self.assertEqual(e.tag, 'html')
+
     # XXX in _elementtree, the constructor of TreeBuilder expects no
     # arguments
     @unittest.expectedFailure

Modules/_elementtree.c

@@ -1803,23 +1803,22 @@ static PyTypeObject Element_Type = {
 typedef struct {
     PyObject_HEAD
 
-    PyObject* root; /* root node (first created node) */
+    PyObject *root; /* root node (first created node) */
 
-    ElementObject* this; /* current node */
-    ElementObject* last; /* most recently created node */
+    ElementObject *this; /* current node */
+    ElementObject *last; /* most recently created node */
 
-    PyObject* data; /* data collector (string or list), or NULL */
+    PyObject *data; /* data collector (string or list), or NULL */
 
-    PyObject* stack; /* element stack */
-    Py_ssize_t index; /* current stack size (0=empty) */
+    PyObject *stack; /* element stack */
+    Py_ssize_t index; /* current stack size (0 means empty) */
 
     /* element tracing */
-    PyObject* events; /* list of events, or NULL if not collecting */
-    PyObject* start_event_obj; /* event objects (NULL to ignore) */
-    PyObject* end_event_obj;
-    PyObject* start_ns_event_obj;
-    PyObject* end_ns_event_obj;
-
+    PyObject *events; /* list of events, or NULL if not collecting */
+    PyObject *start_event_obj; /* event objects (NULL to ignore) */
+    PyObject *end_event_obj;
+    PyObject *start_ns_event_obj;
+    PyObject *end_ns_event_obj;
 } TreeBuilderObject;
 
 static PyTypeObject TreeBuilder_Type;
@@ -1829,48 +1828,42 @@ static PyTypeObject TreeBuilder_Type;
 /* -------------------------------------------------------------------- */
 /* constructor and destructor */
 
-LOCAL(PyObject*)
-treebuilder_new(void)
+static PyObject *
+treebuilder_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 {
-    TreeBuilderObject* self;
-
-    self = PyObject_New(TreeBuilderObject, &TreeBuilder_Type);
-    if (self == NULL)
-        return NULL;
-
-    self->root = NULL;
+    TreeBuilderObject *t = (TreeBuilderObject *)type->tp_alloc(type, 0);
+    if (t != NULL) {
+        t->root = NULL;
 
         Py_INCREF(Py_None);
-    self->this = (ElementObject*) Py_None;
-
+        t->this = (ElementObject *)Py_None;
         Py_INCREF(Py_None);
-    self->last = (ElementObject*) Py_None;
+        t->last = (ElementObject *)Py_None;
 
-    self->data = NULL;
-
-    self->stack = PyList_New(20);
-    self->index = 0;
-
-    self->events = NULL;
-    self->start_event_obj = self->end_event_obj = NULL;
-    self->start_ns_event_obj = self->end_ns_event_obj = NULL;
-
-    ALLOC(sizeof(TreeBuilderObject), "create treebuilder");
+        t->data = NULL;
+        t->stack = PyList_New(20);
+        if (!t->stack) {
+            Py_DECREF(t->this);
+            Py_DECREF(t->last);
+            return NULL;
+        }
+        t->index = 0;
 
-    return (PyObject*) self;
+        t->events = NULL;
+        t->start_event_obj = t->end_event_obj = NULL;
+        t->start_ns_event_obj = t->end_ns_event_obj = NULL;
+    }
+    return (PyObject *)t;
 }
 
-static PyObject*
-treebuilder(PyObject* self_, PyObject* args)
+static int
+treebuilder_init(PyObject *self, PyObject *args, PyObject *kwds)
 {
-    if (!PyArg_ParseTuple(args, ":TreeBuilder"))
-        return NULL;
-
-    return treebuilder_new();
+    return 0;
 }
 
 static void
-treebuilder_dealloc(TreeBuilderObject* self)
+treebuilder_dealloc(TreeBuilderObject *self)
 {
     Py_XDECREF(self->end_ns_event_obj);
     Py_XDECREF(self->start_ns_event_obj);
@@ -1883,9 +1876,7 @@ treebuilder_dealloc(TreeBuilderObject* self)
     Py_DECREF(self->this);
     Py_XDECREF(self->root);
 
-    RELEASE(sizeof(TreeBuilderObject), "destroy treebuilder");
-
-    PyObject_Del(self);
+    Py_TYPE(self)->tp_free((PyObject *)self);
 }
 
 /* -------------------------------------------------------------------- */
@@ -2189,7 +2180,7 @@ static PyTypeObject TreeBuilder_Type = {
     0,                                              /* tp_getattro */
     0,                                              /* tp_setattro */
     0,                                              /* tp_as_buffer */
-    Py_TPFLAGS_DEFAULT, /* tp_flags */
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,       /* tp_flags */
     0,                                              /* tp_doc */
     0,                                              /* tp_traverse */
     0,                                              /* tp_clear */
@@ -2199,6 +2190,16 @@ static PyTypeObject TreeBuilder_Type = {
     0,                                              /* tp_iternext */
     treebuilder_methods,                            /* tp_methods */
     0,                                              /* tp_members */
+    0,                                              /* tp_getset */
+    0,                                              /* tp_base */
+    0,                                              /* tp_dict */
+    0,                                              /* tp_descr_get */
+    0,                                              /* tp_descr_set */
+    0,                                              /* tp_dictoffset */
+    (initproc)treebuilder_init,                     /* tp_init */
+    PyType_GenericAlloc,                            /* tp_alloc */
+    treebuilder_new,                                /* tp_new */
+    0,                                              /* tp_free */
 };
 
 /* ==================================================================== */
@@ -2684,7 +2685,7 @@ xmlparser(PyObject* self_, PyObject* args, PyObject* kw)
 
     /* setup target handlers */
     if (!target) {
-        target = treebuilder_new();
+        target = treebuilder_new(&TreeBuilder_Type, NULL, NULL);
         if (!target) {
             EXPAT(ParserFree)(self->parser);
             PyObject_Del(self->names);
@@ -3073,7 +3074,6 @@ static PyTypeObject XMLParser_Type = {
 
 static PyMethodDef _functions[] = {
     {"SubElement", (PyCFunction) subelement, METH_VARARGS|METH_KEYWORDS},
-    {"TreeBuilder", (PyCFunction) treebuilder, METH_VARARGS},
 #if defined(USE_EXPAT)
     {"XMLParser", (PyCFunction) xmlparser, METH_VARARGS|METH_KEYWORDS},
 #endif
@@ -3185,5 +3185,8 @@ PyInit__elementtree(void)
     Py_INCREF((PyObject *)&Element_Type);
     PyModule_AddObject(m, "Element", (PyObject *)&Element_Type);
 
+    Py_INCREF((PyObject *)&TreeBuilder_Type);
+    PyModule_AddObject(m, "TreeBuilder", (PyObject *)&TreeBuilder_Type);
+
     return m;
 }