pythongh-115775: Use static analysis to estimate how many items a cla…

…ss dict can contain

Include/internal/pycore_global_strings.h

@@ -112,6 +112,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(__enter__)
         STRUCT_FOR_ID(__eq__)
         STRUCT_FOR_ID(__exit__)
+        STRUCT_FOR_ID(__expected_attributes__)
         STRUCT_FOR_ID(__file__)
         STRUCT_FOR_ID(__float__)
         STRUCT_FOR_ID(__floordiv__)

Lib/test/test_compile.py

@@ -1960,6 +1960,64 @@ def test_load_super_attr(self):
         )
 
 
+class TestExpectedAttributes(unittest.TestCase):
+
+    def test_basic(self):
+        class C:
+            def f(self):
+                self.a = self.b = 42
+
+        self.assertIsInstance(C.__expected_attributes__, tuple)
+        self.assertEqual(sorted(C.__expected_attributes__), ['a', 'b'])
+
+    def test_nested_function(self):
+        class C:
+            def f(self):
+                self.x = 1
+                self.y = 2
+                self.x = 3
+
+            def g(self, obj):
+                self.y = 4
+                self.z = 5
+
+                def h(self, a):
+                    self.u = 6
+                    self.v = 7
+
+                obj.self = 8
+
+        self.assertEqual(sorted(C.__expected_attributes__), ['u', 'v', 'x', 'y', 'z'])
+
+    def test_nested_class(self):
+        class C:
+            def f(self):
+                self.x = 42
+                self.y = 42
+
+            class D:
+                def g(self):
+                    self.y = 42
+                    self.z = 42
+
+        self.assertEqual(sorted(C.__expected_attributes__), ['x', 'y'])
+        self.assertEqual(sorted(C.D.__expected_attributes__), ['y', 'z'])
+
+    def test_subclass(self):
+        class C:
+            def f(self):
+                self.x = 42
+                self.y = 42
+
+        class D(C):
+            def g(self):
+                self.y = 42
+                self.z = 42
+
+        self.assertEqual(sorted(C.__expected_attributes__), ['x', 'y'])
+        self.assertEqual(sorted(D.__expected_attributes__), ['y', 'z'])
+
+
 class TestExpressionStackSize(unittest.TestCase):
     # These tests check that the computed stack size for a code object
     # stays within reasonable bounds (see issue #21523 for an example

Misc/NEWS.d/next/Core and Builtins/2024-02-25-14-17-25.gh-issue-115775.CNbGbJ.rst

@@ -0,0 +1,3 @@
+Compiler populate the new ``__expected_attributes__`` field on a class with
+the names of attributes of this class which are accessed through self.X from
+any function in its body.

Python/compile.c

@@ -359,6 +359,7 @@ struct compiler_unit {
     int u_scope_type;
 
     PyObject *u_private;        /* for private name mangling */
+    PyObject *u_expected_attributes;  /* for class: attributes accessed via self.X */
 
     instr_sequence u_instr_sequence; /* codegen output */
 
@@ -690,9 +691,27 @@ compiler_unit_free(struct compiler_unit *u)
     Py_CLEAR(u->u_metadata.u_cellvars);
     Py_CLEAR(u->u_metadata.u_fasthidden);
     Py_CLEAR(u->u_private);
+    Py_CLEAR(u->u_expected_attributes);
     PyMem_Free(u);
 }
 
+static struct compiler_unit *
+get_class_compiler_unit(struct compiler *c)
+{
+    Py_ssize_t stack_size = PyList_GET_SIZE(c->c_stack);
+    assert(stack_size >= 1);
+    for (Py_ssize_t i = stack_size - 1; i >= 0; i--) {
+        PyObject *capsule = PyList_GET_ITEM(c->c_stack, i);
+        struct compiler_unit *u = (struct compiler_unit *)PyCapsule_GetPointer(
+                                                              capsule, CAPSULE_NAME);
+        assert(u);
+        if (u->u_scope_type == COMPILER_SCOPE_CLASS) {
+            return u;
+        }
+    }
+    return NULL;
+}
+
 static int
 compiler_set_qualname(struct compiler *c)
 {
@@ -1336,6 +1355,16 @@ compiler_enter_scope(struct compiler *c, identifier name,
     }
 
     u->u_private = NULL;
+    if (scope_type == COMPILER_SCOPE_CLASS) {
+        u->u_expected_attributes = PySet_New(0);
+        if (!u->u_expected_attributes) {
+            compiler_unit_free(u);
+            return ERROR;
+        }
+    }
+    else {
+        u->u_expected_attributes = NULL;
+    }
 
     /* Push the old compiler_unit on the stack. */
     if (c->u) {
@@ -2517,6 +2546,17 @@ compiler_class_body(struct compiler *c, stmt_ty s, int firstlineno)
         compiler_exit_scope(c);
         return ERROR;
     }
+    assert(c->u->u_expected_attributes);
+    PyObject *expected_attributes = PySequence_Tuple(c->u->u_expected_attributes);
+    if (expected_attributes == NULL) {
+        compiler_exit_scope(c);
+        return ERROR;
+    }
+    ADDOP_LOAD_CONST(c, NO_LOCATION, expected_attributes);
+    if (compiler_nameop(c, NO_LOCATION, &_Py_ID(__expected_attributes__), Store) < 0) {
+        compiler_exit_scope(c);
+        return ERROR;
+    }
     /* The following code is artificial */
     /* Set __classdictcell__ if necessary */
     if (c->u->u_ste->ste_needs_classdict) {
@@ -2657,6 +2697,7 @@ compiler_class(struct compiler *c, stmt_ty s)
                                                          s->v.ClassDef.keywords));
 
         PyCodeObject *co = optimize_and_assemble(c, 0);
+
         compiler_exit_scope(c);
         if (co == NULL) {
             return ERROR;
@@ -6246,6 +6287,17 @@ compiler_visit_expr1(struct compiler *c, expr_ty e)
             ADDOP(c, loc, NOP);
             return SUCCESS;
         }
+        if (e->v.Attribute.value->kind == Name_kind &&
+            _PyUnicode_EqualToASCIIString(e->v.Attribute.value->v.Name.id, "self"))
+        {
+            struct compiler_unit *class_u = get_class_compiler_unit(c);
+            if (class_u != NULL) {
+                assert(class_u->u_scope_type == COMPILER_SCOPE_CLASS);
+                assert(class_u->u_expected_attributes);
+                RETURN_IF_ERROR(
+                    PySet_Add(class_u->u_expected_attributes, e->v.Attribute.attr));
+            }
+        }
         VISIT(c, expr, e->v.Attribute.value);
         loc = LOC(e);
         loc = update_start_location_to_match_attr(c, loc, e);