gh-124697: Represent inlined comprehensions as subscopes in the symbol table

Author: iritkatriel

Doc/library/symtable.rst

@@ -57,6 +57,14 @@ Examining Symbol Tables
 
       Used for the symbol table of a class.
 
+   .. attribute:: INLINED_COMPREHENSION
+      :value: "inlined comprehension"
+
+      Used for the symbol table of a list, set or dict comprehension that
+      is inlined into the enclosing code unit (see :pep:`709`). A symbol
+      table of this type represents a sub-scope of the enclosing code unit's
+      scope, and it does not correspond to a separate compilation unit.
+
    The following members refer to different flavors of
    :ref:`annotation scopes <annotation-scopes>`.
 

Include/internal/pycore_compile.h

@@ -69,9 +69,8 @@ typedef struct {
     PyObject *u_varnames;  /* local variables */
     PyObject *u_cellvars;  /* cell variables */
     PyObject *u_freevars;  /* free variables */
-    PyObject *u_fasthidden; /* dict; keys are names that are fast-locals only
-                               temporarily within an inlined comprehension. When
-                               value is True, treat as fast-local. */
+    PyObject *u_fasthidden; /* set of names that are fast-locals only
+                               temporarily within an inlined comprehension. */
 
     Py_ssize_t u_argcount;        /* number of arguments for block */
     Py_ssize_t u_posonlyargcount;        /* number of positional only arguments for block */
@@ -155,7 +154,6 @@ int _PyCompile_ResolveNameop(struct _PyCompiler *c, PyObject *mangled, int scope
                              _PyCompile_optype *optype, Py_ssize_t *arg);
 
 int _PyCompile_IsInteractiveTopLevel(struct _PyCompiler *c);
-int _PyCompile_IsInInlinedComp(struct _PyCompiler *c);
 int _PyCompile_ScopeType(struct _PyCompiler *c);
 int _PyCompile_OptimizationLevel(struct _PyCompiler *c);
 int _PyCompile_LookupArg(struct _PyCompiler *c, PyCodeObject *co, PyObject *name);
@@ -179,16 +177,15 @@ enum {
 
 typedef struct {
     PyObject *pushed_locals;
-    PyObject *temp_symbols;
-    PyObject *fast_hidden;
     _PyJumpTargetLabel cleanup;
+    PySTEntryObject *saved_ste;
 } _PyCompile_InlinedComprehensionState;
 
-int _PyCompile_TweakInlinedComprehensionScopes(struct _PyCompiler *c, _Py_SourceLocation loc,
-                                               PySTEntryObject *entry,
-                                               _PyCompile_InlinedComprehensionState *state);
-int _PyCompile_RevertInlinedComprehensionScopes(struct _PyCompiler *c, _Py_SourceLocation loc,
-                                                _PyCompile_InlinedComprehensionState *state);
+int _PyCompile_EnterInlinedComprehensionScope(struct _PyCompiler *c, _Py_SourceLocation loc,
+                                              PySTEntryObject *entry,
+                                              _PyCompile_InlinedComprehensionState *state);
+int _PyCompile_ExitInlinedComprehensionScope(struct _PyCompiler *c, _Py_SourceLocation loc,
+                                             _PyCompile_InlinedComprehensionState *state);
 int _PyCompile_AddDeferredAnnotation(struct _PyCompiler *c, stmt_ty s,
                                      PyObject **conditional_annotation_index);
 void _PyCompile_EnterConditionalBlock(struct _PyCompiler *c);

Include/internal/pycore_symtable.h

@@ -33,6 +33,10 @@ typedef enum _block_type {
     // i.e., a TypeVar, a TypeVarTuple or a ParamSpec object (the latter two
     // do not support a bound or a constraint tuple).
     TypeVariableBlock,
+    // Comprehension which is inlined into the enclosing code unit (see PEP 709).
+    // Represents a sub-scope of the enclosing code unit's scope rather than a
+    // separate scope.
+    InlinedComprehensionBlock,
 } _Py_block_ty;
 
 typedef enum _comprehension_type {
@@ -119,7 +123,6 @@ typedef struct _symtable_entry {
                                              should be created */
     unsigned ste_needs_classdict : 1; /* for class scopes, true if a closure
                                          over the class dict should be created */
-    unsigned ste_comp_inlined : 1; /* true if this comprehension is inlined */
     unsigned ste_comp_iter_target : 1; /* true if visiting comprehension target */
     unsigned ste_can_see_class_scope : 1; /* true if this block can see names bound in an
                                              enclosing class scope */
@@ -132,6 +135,7 @@ typedef struct _symtable_entry {
     int ste_comp_iter_expr; /* non-zero if visiting a comprehension range expression */
     _Py_SourceLocation ste_loc; /* source location of block */
     struct _symtable_entry *ste_annotation_block; /* symbol table entry for this entry's annotations */
+    struct _symtable_entry *ste_parent; /* st entry for the enclosing block if this entry is a sub-scope, NULL otherwise */
     struct symtable *ste_table;
 } PySTEntryObject;
 

Lib/symtable.py

@@ -56,6 +56,7 @@ class SymbolTableType(StrEnum):
     TYPE_ALIAS = "type alias"
     TYPE_PARAMETERS = "type parameters"
     TYPE_VARIABLE = "type variable"
+    INLINED_COMPREHENSION = "inlined comprehension"
 
 
 class SymbolTable:
@@ -98,6 +99,8 @@ def get_type(self):
             return SymbolTableType.TYPE_PARAMETERS
         if self._table.type == _symtable.TYPE_TYPE_VARIABLE:
             return SymbolTableType.TYPE_VARIABLE
+        if self._table.type == _symtable.TYPE_INLINED_COMPREHENSION:
+            return SymbolTableType.INLINED_COMPREHENSION
         assert False, f"unexpected type: {self._table.type}"
 
     def get_id(self):

Lib/test/test_compiler_assemble.py

@@ -17,8 +17,9 @@ def complete_metadata(self, metadata, filename="myfile.py"):
             metadata.setdefault(key, key)
         for key in ['consts']:
             metadata.setdefault(key, [])
-        for key in ['names', 'varnames', 'cellvars', 'freevars', 'fasthidden']:
+        for key in ['names', 'varnames', 'cellvars', 'freevars']:
             metadata.setdefault(key, {})
+        metadata.setdefault('fasthidden', None)
         for key in ['argcount', 'posonlyargcount', 'kwonlyargcount']:
             metadata.setdefault(key, 0)
         metadata.setdefault('firstlineno', 1)

Lib/test/test_symtable.py

@@ -426,11 +426,11 @@ def test_symbol_repr(self):
                          "<symbol 'T': LOCAL, DEF_LOCAL|DEF_TYPE_PARAM>")
 
         st1 = symtable.symtable("[x for x in [1]]", "?", "exec")
-        self.assertEqual(repr(st1.lookup("x")),
+        self.assertEqual(repr(st1.get_children()[0].lookup("x")),
                          "<symbol 'x': LOCAL, USE|DEF_LOCAL|DEF_COMP_ITER>")
 
         st2 = symtable.symtable("[(lambda: x) for x in [1]]", "?", "exec")
-        self.assertEqual(repr(st2.lookup("x")),
+        self.assertEqual(repr(st2.get_children()[0].lookup("x")),
                          "<symbol 'x': CELL, DEF_LOCAL|DEF_COMP_ITER|DEF_COMP_CELL>")
 
         st3 = symtable.symtable("def f():\n"
@@ -502,6 +502,64 @@ def test_nested_genexpr(self):
         self.assertEqual(sorted(st.get_identifiers()), [".0", "y"])
         self.assertEqual(st.get_children(), [])
 
+    def test_inlined_comprehension_in_genexpr(self):
+        st = symtable.symtable("([y for y in x] for x in a)", "?", "exec")
+        self.assertEqual(len(st.get_children()), 1)
+        st = st.get_children()[0]
+        self.assertIs(st.get_type(), symtable.SymbolTableType.FUNCTION)
+        self.assertEqual(st.get_name(), "<genexpr>")
+        self.assertFalse(st.is_nested())
+        self.assertEqual(sorted(st.get_identifiers()), [".0", "x"])
+        children = st.get_children()
+        self.assertEqual(len(children), 1)
+        self.check_inlined_listcomp(children[0], ["y"], nested=True)
+
+    def check_inlined_listcomp(self, st, identifiers, *, nested, nchildren=0):
+        self.assertIs(st.get_type(), symtable.SymbolTableType.INLINED_COMPREHENSION)
+        self.assertEqual(st.get_name(), "<listcomp>")
+        self.assertEqual(st.is_nested(), nested)
+        self.assertEqual(sorted(st.get_identifiers()), identifiers)
+        children = st.get_children()
+        self.assertEqual(len(children), nchildren)
+        return children
+
+    def check_nested_inlined_listcomp(self, outer, hoisted, outer_ids, inner_ids, *, nested):
+        # Nested namespaces of inlined comprehensions are also hoisted into
+        # the enclosing scope's children list.
+        inner, = self.check_inlined_listcomp(
+            outer, outer_ids, nested=nested, nchildren=1)
+        self.check_inlined_listcomp(inner, inner_ids, nested=True)
+        self.assertIs(hoisted, inner)
+
+    def test_inlined_comprehension(self):
+        st = symtable.symtable("[x for x in [1]]", "?", "exec")
+        self.assertEqual(sorted(st.get_identifiers()), [])
+        children = st.get_children()
+        self.assertEqual(len(children), 1)
+        self.check_inlined_listcomp(children[0], ["x"], nested=False)
+
+    def test_inlined_nested_comprehension(self):
+        st = symtable.symtable("[[y for y in x] for x in [1]]", "?", "exec")
+        self.assertEqual(sorted(st.get_identifiers()), [])
+        children = st.get_children()
+        self.assertEqual(len(children), 2)
+        self.check_nested_inlined_listcomp(
+            children[0], children[1], ["x"], ["y"], nested=False)
+
+    def test_inlined_sibling_nested_comprehensions(self):
+        st = symtable.symtable(
+            "def f(): [[y for y in x] for x in [1]]; [[w for w in z] for z in [2]]",
+            "?", "exec")
+        f = find_block(st, "f")
+        self.assertIs(f.get_type(), symtable.SymbolTableType.FUNCTION)
+        self.assertEqual(sorted(f.get_identifiers()), [])
+        children = f.get_children()
+        self.assertEqual(len(children), 4)
+        self.check_nested_inlined_listcomp(
+            children[0], children[1], ["x"], ["y"], nested=True)
+        self.check_nested_inlined_listcomp(
+            children[2], children[3], ["z"], ["w"], nested=True)
+
     def test__symtable_refleak(self):
         # Regression test for reference leak in PyUnicode_FSDecoder.
         # See https://github.com/python/cpython/issues/139748.

Modules/_testinternalcapi.c

@@ -1357,13 +1357,16 @@ _testinternalcapi_assemble_code_object_impl(PyObject *module,
     umd.u_cellvars = PyDict_GetItemString(metadata, "cellvars");
     umd.u_freevars = PyDict_GetItemString(metadata, "freevars");
     umd.u_fasthidden = PyDict_GetItemString(metadata, "fasthidden");
+    if (umd.u_fasthidden == Py_None) {
+        umd.u_fasthidden = NULL;
+    }
 
     assert(PyDict_Check(umd.u_consts));
     assert(PyDict_Check(umd.u_names));
     assert(PyDict_Check(umd.u_varnames));
     assert(PyDict_Check(umd.u_cellvars));
     assert(PyDict_Check(umd.u_freevars));
-    assert(PyDict_Check(umd.u_fasthidden));
+    assert(umd.u_fasthidden == NULL || PySet_Check(umd.u_fasthidden));
 
     umd.u_argcount = get_nonnegative_int_from_dict(metadata, "argcount");
     umd.u_posonlyargcount = get_nonnegative_int_from_dict(metadata, "posonlyargcount");

Modules/symtablemodule.c

@@ -144,6 +144,8 @@ symtable_init_constants(PyObject *m)
         return -1;
     if (PyModule_AddIntConstant(m, "TYPE_TYPE_VARIABLE", TypeVariableBlock) < 0)
         return -1;
+    if (PyModule_AddIntConstant(m, "TYPE_INLINED_COMPREHENSION", InlinedComprehensionBlock) < 0)
+        return -1;
 
     if (PyModule_AddIntMacro(m, LOCAL) < 0) return -1;
     if (PyModule_AddIntMacro(m, GLOBAL_EXPLICIT) < 0) return -1;

Python/assemble.c

@@ -520,13 +520,15 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus,
 
             _PyLocals_Kind kind = CO_FAST_LOCAL | argvarkinds[i].kind;
 
-            int has_key = PyDict_Contains(umd->u_fasthidden, k);
-            RETURN_IF_ERROR(has_key);
-            if (has_key) {
-                kind |= CO_FAST_HIDDEN;
+            if (umd->u_fasthidden != NULL) {
+                int hidden = PySet_Contains(umd->u_fasthidden, k);
+                RETURN_IF_ERROR(hidden);
+                if (hidden) {
+                    kind |= CO_FAST_HIDDEN;
+                }
             }
 
-            has_key = PyDict_Contains(umd->u_cellvars, k);
+            int has_key = PyDict_Contains(umd->u_cellvars, k);
             RETURN_IF_ERROR(has_key);
             if (has_key) {
                 kind |= CO_FAST_CELL;

Python/codegen.c

@@ -3342,7 +3342,7 @@ codegen_nameop(compiler *c, location loc,
     case COMPILE_OP_DEREF:
         switch (ctx) {
         case Load:
-            if (SYMTABLE_ENTRY(c)->ste_type == ClassBlock && !_PyCompile_IsInInlinedComp(c)) {
+            if (SYMTABLE_ENTRY(c)->ste_type == ClassBlock) {
                 op = LOAD_FROM_DICT_OR_DEREF;
                 // First load the locals
                 if (codegen_addop_noarg(INSTR_SEQUENCE(c), LOAD_LOCALS, loc) < 0) {
@@ -3395,8 +3395,9 @@ codegen_nameop(compiler *c, location loc,
     case COMPILE_OP_NAME:
         switch (ctx) {
         case Load:
-            op = (SYMTABLE_ENTRY(c)->ste_type == ClassBlock
-                    && _PyCompile_IsInInlinedComp(c))
+            /* LOAD_NAME in a class reads the class dict; inlined comps must not. */
+            op = (SCOPE_TYPE(c) == COMPILE_SCOPE_CLASS
+                  && SYMTABLE_ENTRY(c)->ste_type == InlinedComprehensionBlock)
                 ? LOAD_GLOBAL
                 : LOAD_NAME;
             break;
@@ -4917,9 +4918,6 @@ codegen_push_inlined_comprehension_locals(compiler *c, location loc,
                                           PySTEntryObject *comp,
                                           _PyCompile_InlinedComprehensionState *state)
 {
-    int in_class_block = (SYMTABLE_ENTRY(c)->ste_type == ClassBlock) &&
-                          !_PyCompile_IsInInlinedComp(c);
-    PySTEntryObject *outer = SYMTABLE_ENTRY(c);
     // iterate over names bound in the comprehension and ensure we isolate
     // them from the outer scope as needed
     PyObject *k, *v;
@@ -4930,11 +4928,7 @@ codegen_push_inlined_comprehension_locals(compiler *c, location loc,
         RETURN_IF_ERROR(symbol);
         long scope = SYMBOL_TO_SCOPE(symbol);
 
-        long outsymbol = _PyST_GetSymbol(outer, k);
-        RETURN_IF_ERROR(outsymbol);
-        long outsc = SYMBOL_TO_SCOPE(outsymbol);
-
-        if ((symbol & DEF_LOCAL && !(symbol & DEF_NONLOCAL)) || in_class_block) {
+        if ((symbol & DEF_LOCAL) && !(symbol & DEF_NONLOCAL)) {
             // local names bound in comprehension must be isolated from
             // outer scope; push existing value (which may be NULL if
             // not defined) on stack
@@ -4949,15 +4943,17 @@ codegen_push_inlined_comprehension_locals(compiler *c, location loc,
             // comprehension and restore the original one after
             ADDOP_NAME(c, loc, LOAD_FAST_AND_CLEAR, k, varnames);
             if (scope == CELL) {
-                if (outsc == FREE) {
-                    ADDOP_NAME(c, loc, MAKE_CELL, k, freevars);
-                } else {
-                    ADDOP_NAME(c, loc, MAKE_CELL, k, cellvars);
-                }
+                ADDOP_NAME(c, loc, MAKE_CELL, k, cellvars);
             }
             if (PyList_Append(state->pushed_locals, k) < 0) {
                 return ERROR;
             }
+            if (METADATA(c)->u_fasthidden != NULL) {
+                /* For Module/Class scopes, assemble needs to set CO_FAST_HIDDEN on these names */
+                if (PySet_Add(METADATA(c)->u_fasthidden, k) < 0) {
+                    return ERROR;
+                }
+            }
         }
     }
     if (state->pushed_locals) {
@@ -4986,7 +4982,7 @@ push_inlined_comprehension_state(compiler *c, location loc,
                                  _PyCompile_InlinedComprehensionState *state)
 {
     RETURN_IF_ERROR(
-        _PyCompile_TweakInlinedComprehensionScopes(c, loc, comp, state));
+        _PyCompile_EnterInlinedComprehensionScope(c, loc, comp, state));
     RETURN_IF_ERROR(
         codegen_push_inlined_comprehension_locals(c, loc, comp, state));
     return SUCCESS;
@@ -5044,7 +5040,7 @@ pop_inlined_comprehension_state(compiler *c, location loc,
                                 _PyCompile_InlinedComprehensionState *state)
 {
     RETURN_IF_ERROR(codegen_pop_inlined_comprehension_locals(c, loc, state));
-    RETURN_IF_ERROR(_PyCompile_RevertInlinedComprehensionScopes(c, loc, state));
+    RETURN_IF_ERROR(_PyCompile_ExitInlinedComprehensionScope(c, loc, state));
     return SUCCESS;
 }
 
@@ -5054,13 +5050,13 @@ codegen_comprehension(compiler *c, expr_ty e, int type,
                       expr_ty val, bool avoid_creation)
 {
     PyCodeObject *co = NULL;
-    _PyCompile_InlinedComprehensionState inline_state = {NULL, NULL, NULL, NO_LABEL};
+    _PyCompile_InlinedComprehensionState inline_state = {NULL, NO_LABEL, NULL};
     comprehension_ty outermost;
     PySTEntryObject *entry = _PySymtable_Lookup(SYMTABLE(c), (void *)e);
     if (entry == NULL) {
         goto error;
     }
-    int is_inlined = entry->ste_comp_inlined;
+    int is_inlined = (entry->ste_type == InlinedComprehensionBlock);
     int is_async_comprehension = entry->ste_coroutine;
 
     location loc = LOC(e);
@@ -5069,7 +5065,7 @@ codegen_comprehension(compiler *c, expr_ty e, int type,
     IterStackPosition iter_state;
     if (is_inlined) {
         VISIT(c, expr, outermost->iter);
-        if (push_inlined_comprehension_state(c, loc, entry, &inline_state)) {
+        if (push_inlined_comprehension_state(c, loc, entry, &inline_state) < 0) {
             goto error;
         }
         iter_state = ITERABLE_ON_STACK;
@@ -5140,8 +5136,8 @@ codegen_comprehension(compiler *c, expr_ty e, int type,
     }
 
     if (is_inlined) {
-        if (pop_inlined_comprehension_state(c, loc, &inline_state)) {
-            goto error;
+        if (pop_inlined_comprehension_state(c, loc, &inline_state) < 0) {
+            goto error_in_scope;
         }
         return SUCCESS;
     }
@@ -5181,15 +5177,18 @@ codegen_comprehension(compiler *c, expr_ty e, int type,
 
     return SUCCESS;
 error_in_scope:
-    if (!is_inlined) {
+    if (is_inlined) {
+        if (inline_state.saved_ste != NULL) {
+            pop_inlined_comprehension_state(c, loc, &inline_state);
+        }
+    }
+    else {
         _PyCompile_ExitScope(c);
     }
 error:
     Py_XDECREF(co);
     Py_XDECREF(entry);
     Py_XDECREF(inline_state.pushed_locals);
-    Py_XDECREF(inline_state.temp_symbols);
-    Py_XDECREF(inline_state.fast_hidden);
     return ERROR;
 }