Closes #15512: Correct __sizeof__ support for parser

5 files changed, 103 insertions, 4 deletions

diff --git a/Include/node.h b/Include/node.h
index e23e709..9f6760c 100644
--- a/Include/node.h
+++ b/Include/node.h
@@ -20,6 +20,9 @@ PyAPI_FUNC(node *) PyNode_New(int type);
 PyAPI_FUNC(int) PyNode_AddChild(node *n, int type,
                                       char *str, int lineno, int col_offset);
 PyAPI_FUNC(void) PyNode_Free(node *n);
+#ifndef Py_LIMITED_API
+Py_ssize_t _PyNode_SizeOf(node *n);
+#endif
 
 /* Node access functions */
 #define NCH(n)		((n)->n_nchildren)
diff --git a/Lib/test/test_parser.py b/Lib/test/test_parser.py
index 33b91bd..c6900de 100644
--- a/Lib/test/test_parser.py
+++ b/Lib/test/test_parser.py
@@ -1,7 +1,8 @@
 import parser
 import unittest
 import sys
-from test import test_support
+import struct
+from test import test_support as support
 
 #
 #  First, we test that we can generate trees from valid source fragments,
@@ -583,12 +584,59 @@ class ParserStackLimitTestCase(unittest.TestCase):
         print >>sys.stderr, "Expecting 's_push: parser stack overflow' in next line"
         self.assertRaises(MemoryError, parser.expr, e)
 
+class STObjectTestCase(unittest.TestCase):
+    """Test operations on ST objects themselves"""
+
+    check_sizeof = support.check_sizeof
+
+    @support.cpython_only
+    def test_sizeof(self):
+        def XXXROUNDUP(n):
+            if n <= 1:
+                return n
+            if n <= 128:
+                return (n + 3) & ~3
+            return 1 << (n - 1).bit_length()
+
+        basesize = support.calcobjsize('Pii')
+        nodesize = struct.calcsize('hP3iP0h')
+        def sizeofchildren(node):
+            if node is None:
+                return 0
+            res = 0
+            hasstr = len(node) > 1 and isinstance(node[-1], str)
+            if hasstr:
+                res += len(node[-1]) + 1
+            children = node[1:-1] if hasstr else node[1:]
+            if children:
+                res += XXXROUNDUP(len(children)) * nodesize
+            res1 = res
+            if children:
+                for child in children:
+                    res += sizeofchildren(child)
+            return res
+
+        def check_st_sizeof(st):
+            self.check_sizeof(st, basesize + nodesize +
+                                  sizeofchildren(st.totuple()))
+
+        check_st_sizeof(parser.expr('2 + 3'))
+        check_st_sizeof(parser.expr('2 + 3 + 4'))
+        check_st_sizeof(parser.suite('x = 2 + 3'))
+        check_st_sizeof(parser.suite(''))
+        check_st_sizeof(parser.suite('# -*- coding: utf-8 -*-'))
+        check_st_sizeof(parser.expr('[' + '2,' * 1000 + ']'))
+
+
+    # XXX tests for pickling and unpickling of ST objects should go here
+
 def test_main():
-    test_support.run_unittest(
+    support.run_unittest(
         RoundtripLegalSyntaxTestCase,
         IllegalSyntaxTestCase,
         CompileTestCase,
         ParserStackLimitTestCase,
+        STObjectTestCase,
     )
 
 
diff --git a/Misc/NEWS b/Misc/NEWS
index 635f847..b6e4791 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -107,6 +107,9 @@ Library
 - Issue #15487: Add a __sizeof__ implementation for buffered I/O objects.
   Patch by Serhiy Storchaka.
 
+- Issue #15512: Add a __sizeof__ implementation for parser.
+  Patch by Serhiy Storchaka.
+
 - Issue #15402: An issue in the struct module that caused sys.getsizeof to
   return incorrect results for struct.Struct instances has been fixed.
   Initial patch by Serhiy Storchaka.
diff --git a/Modules/parsermodule.c b/Modules/parsermodule.c
index 632475c..a7cdc8f 100644
--- a/Modules/parsermodule.c
+++ b/Modules/parsermodule.c
@@ -169,8 +169,10 @@ typedef struct {
 
 
 static void parser_free(PyST_Object *st);
+static PyObject* parser_sizeof(PyST_Object *, void *);
 static int parser_compare(PyST_Object *left, PyST_Object *right);
 static PyObject *parser_getattr(PyObject *self, char *name);
+static PyMethodDef parser_methods[];
 
 
 static
@@ -200,7 +202,14 @@ PyTypeObject PyST_Type = {
     Py_TPFLAGS_DEFAULT,                 /* tp_flags             */
 
     /* __doc__ */
-    "Intermediate representation of a Python parse tree."
+    "Intermediate representation of a Python parse tree.",
+    0,                                  /* tp_traverse */
+    0,                                  /* tp_clear */
+    0,                                  /* tp_richcompare */
+    0,                                  /* tp_weaklistoffset */
+    0,                                  /* tp_iter */
+    0,                                  /* tp_iternext */
+    parser_methods,                     /* tp_methods */
 };  /* PyST_Type */
 
 
@@ -508,7 +517,8 @@ parser_methods[] = {
         PyDoc_STR("Creates a list-tree representation of this ST.")},
     {"totuple",         (PyCFunction)parser_st2tuple,   PUBLIC_METHOD_TYPE,
         PyDoc_STR("Creates a tuple-tree representation of this ST.")},
-
+    {"__sizeof__",      (PyCFunction)parser_sizeof,     METH_NOARGS,
+        PyDoc_STR("Returns size in memory, in bytes.")},
     {NULL, NULL, 0, NULL}
 };
 
@@ -695,6 +705,15 @@ parser_tuple2ast(PyST_Object *self, PyObject *args, PyObject *kw)
     return parser_tuple2st(self, args, kw);
 }
 
+static PyObject *
+parser_sizeof(PyST_Object *st, void *unused)
+{
+    Py_ssize_t res;
+
+    res = sizeof(PyST_Object) + _PyNode_SizeOf(st->st_node);
+    return PyLong_FromSsize_t(res);
+}
+
 
 /*  node* build_node_children()
  *
diff --git a/Parser/node.c b/Parser/node.c
index 9eba76b..0dea30f 100644
--- a/Parser/node.c
+++ b/Parser/node.c
@@ -114,6 +114,7 @@ PyNode_AddChild(register node *n1, int type, char *str, int lineno, int col_offs
 
 /* Forward */
 static void freechildren(node *);
+static Py_ssize_t sizeofchildren(node *n);
 
 
 void
@@ -125,6 +126,16 @@ PyNode_Free(node *n)
     }
 }
 
+Py_ssize_t
+_PyNode_SizeOf(node *n)
+{
+    Py_ssize_t res = 0;
+
+    if (n != NULL)
+        res = sizeof(node) + sizeofchildren(n);
+    return res;
+}
+
 static void
 freechildren(node *n)
 {
@@ -136,3 +147,18 @@ freechildren(node *n)
     if (STR(n) != NULL)
         PyObject_FREE(STR(n));
 }
+
+static Py_ssize_t
+sizeofchildren(node *n)
+{
+    Py_ssize_t res = 0;
+    int i;
+    for (i = NCH(n); --i >= 0; )
+        res += sizeofchildren(CHILD(n, i));
+    if (n->n_child != NULL)
+        /* allocated size of n->n_child array */
+        res += XXXROUNDUP(NCH(n)) * sizeof(node);
+    if (STR(n) != NULL)
+        res += strlen(STR(n)) + 1;
+    return res;
+}