Issue #14055: Add __sizeof__ support to _elementtree.

diff --git a/Lib/test/test_xml_etree_c.py b/Lib/test/test_xml_etree_c.py
index 142a22ff980c93ff4fdf637f4a2ce97aa342c834..65b40d89f1f8cb93d76d0ea6e6ab14381005fcc6 100644 (file)
--- 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 4d221d5b40c115083dfba737d5c3d83e3dbf7887..18e381b6b9f2ea6e98a42b22c2ce5bfdc7c844a9 100644 (file)
--- 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 cb840485809accc3f8f9d0b1c3df5ad0e87f1acd..6415797f5690a6e7e384b98e2783c2335e7aa079 100644 (file)
--- 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}
 };