Issue #14055: Add __sizeof__ support to _elementtree.

3 files changed, 53 insertions, 2 deletions

diff --git a/Lib/test/test_xml_etree_c.py b/Lib/test/test_xml_etree_c.py
index 142a22f..65b40d8 100644
--- a/Lib/test/test_xml_etree_c.py
+++ b/Lib/test/test_xml_etree_c.py
@@ -1,5 +1,5 @@
 # xml.etree test for cElementTree
-
+import sys, struct
 from test import support
 from test.support import import_fresh_module
 import unittest
@@ -40,6 +40,40 @@ class TestAcceleratorImported(unittest.TestCase):
         self.assertEqual(cET_alias.SubElement.__module__, '_elementtree')
 
 
+@unittest.skipUnless(cET, 'requires _elementtree')
+class SizeofTest(unittest.TestCase):
+    def setUp(self):
+        import _testcapi
+        gc_headsize = _testcapi.SIZEOF_PYGC_HEAD
+        # object header
+        header = 'PP'
+        if hasattr(sys, "gettotalrefcount"):
+            # debug header
+            header = 'PP' + header
+        # fields
+        element = header + '5P'
+        self.elementsize = gc_headsize + struct.calcsize(element)
+        # extra
+        self.extra = struct.calcsize('PiiP4P')
+
+    def test_element(self):
+        e = cET.Element('a')
+        self.assertEqual(sys.getsizeof(e), self.elementsize)
+
+    def test_element_with_attrib(self):
+        e = cET.Element('a', href='about:')
+        self.assertEqual(sys.getsizeof(e), 
+                         self.elementsize + self.extra)
+
+    def test_element_with_children(self):
+        e = cET.Element('a')
+        for i in range(5):
+            cET.SubElement(e, 'span')
+        # should have space for 8 children now
+        self.assertEqual(sys.getsizeof(e), 
+                         self.elementsize + self.extra + 
+                         struct.calcsize('8P'))
+
 def test_main():
     from test import test_xml_etree, test_xml_etree_c
 
@@ -47,7 +81,8 @@ def test_main():
     support.run_unittest(
         MiscTests,
         TestAliasWorking,
-        TestAcceleratorImported
+        TestAcceleratorImported,
+        SizeofTest,
         )
 
     # Run the same test suite as the Python module
diff --git a/Misc/NEWS b/Misc/NEWS
index 4d221d5..18e381b 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -29,6 +29,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #14055: Add __sizeof__ support to _elementtree.
+
 - Issue #15054: A bug in tokenize.tokenize that caused string literals
   with 'b' prefixes to be incorrectly tokenized has been fixed.
   Patch by Serhiy Storchaka.
diff --git a/Modules/_elementtree.c b/Modules/_elementtree.c
index cb84048..6415797 100644
--- a/Modules/_elementtree.c
+++ b/Modules/_elementtree.c
@@ -842,6 +842,19 @@ element_deepcopy(ElementObject* self, PyObject* args)
     return NULL;
 }
 
+static PyObject*
+element_sizeof(PyObject* _self, PyObject* args)
+{
+    ElementObject *self = (ElementObject*)_self;
+    Py_ssize_t result = sizeof(ElementObject);
+    if (self->extra) {
+        result += sizeof(ElementObjectExtra);
+        if (self->extra->children != self->extra->_children)
+            result += sizeof(PyObject*) * self->extra->allocated;
+    }
+    return PyLong_FromSsize_t(result);
+}
+
 LOCAL(int)
 checkpath(PyObject* tag)
 {
@@ -1609,6 +1622,7 @@ static PyMethodDef element_methods[] = {
 
     {"__copy__", (PyCFunction) element_copy, METH_VARARGS},
     {"__deepcopy__", (PyCFunction) element_deepcopy, METH_VARARGS},
+    {"__sizeof__", element_sizeof, METH_NOARGS},
 
     {NULL, NULL}
 };