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

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>`.
 

Doc/whatsnew/3.16.rst

@@ -498,6 +498,14 @@ symtable
   like the builtin :func:`compile`.
   (Contributed by Serhiy Storchaka in :gh:`153844`.)
 
+* Inlined list, set and dict comprehensions (:pep:`709`) are now represented
+  as their own symbol table entries, of type
+  :attr:`~symtable.SymbolTableType.INLINED_COMPREHENSION`. Each such entry is
+  a lexical child of the enclosing scope and records the comprehension's own
+  locals, cells, and free names. It does not correspond to a separate
+  compilation unit.
+  (Contributed by Irit Katriel in :gh:`124697`.)
+
 
 tkinter
 -------

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,
+                                              PySTEntryObject *entry,
+                                              _PyCompile_InlinedComprehensionState *state);
+int _PyCompile_ExitInlinedComprehensionScope(struct _PyCompiler *c,
+                                             _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;
 
@@ -142,6 +146,7 @@ extern PyTypeObject PySTEntry_Type;
 extern long _PyST_GetSymbol(PySTEntryObject *, PyObject *);
 extern int _PyST_GetScope(PySTEntryObject *, PyObject *);
 extern int _PyST_IsFunctionLike(PySTEntryObject *);
+extern int _PyST_IsClassClosureName(PyObject *);
 
 extern struct symtable* _PySymtable_Build(
     struct _mod *mod,
@@ -172,7 +177,6 @@ _Py_IsPrivateName(PyObject *);
 #define DEF_ANNOT (2<<7)         /* this name is annotated */
 #define DEF_COMP_ITER (2<<8)     /* this name is a comprehension iteration variable */
 #define DEF_TYPE_PARAM (2<<9)    /* this name is a type parameter */
-#define DEF_COMP_CELL (2<<10)    /* this name is a cell in an inlined comprehension */
 
 #define DEF_BOUND (DEF_LOCAL | DEF_PARAM | DEF_IMPORT)
 

Lib/symtable.py

@@ -7,7 +7,7 @@
     DEF_NONLOCAL, DEF_LOCAL,
     DEF_PARAM, DEF_TYPE_PARAM, DEF_FREE_CLASS,
     DEF_IMPORT, DEF_BOUND, DEF_ANNOT,
-    DEF_COMP_ITER, DEF_COMP_CELL,
+    DEF_COMP_ITER,
     SCOPE_OFF, SCOPE_MASK,
     FREE, LOCAL, GLOBAL_IMPLICIT, GLOBAL_EXPLICIT, CELL
 )
@@ -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):
@@ -151,8 +154,10 @@ def lookup(self, name):
             flags = self._table.symbols[name]
             namespaces = self.__check_children(name)
             module_scope = (self._table.name == "top")
+            inlined = (self._table.type == _symtable.TYPE_INLINED_COMPREHENSION)
             sym = self._symbols[name] = Symbol(name, flags, namespaces,
-                                               module_scope=module_scope)
+                                               module_scope=module_scope,
+                                               inlined_comprehension=inlined)
         return sym
 
     def get_symbols(self):
@@ -246,12 +251,14 @@ class Class(SymbolTable):
 
 class Symbol:
 
-    def __init__(self, name, flags, namespaces=None, *, module_scope=False):
+    def __init__(self, name, flags, namespaces=None, *, module_scope=False,
+                 inlined_comprehension=False):
         self.__name = name
         self.__flags = flags
         self.__scope = _get_scope(flags)
         self.__namespaces = namespaces or ()
         self.__module_scope = module_scope
+        self.__inlined_comprehension = inlined_comprehension
 
     def __repr__(self):
         flags_str = '|'.join(self._flags_str())
@@ -345,7 +352,7 @@ def is_comp_iter(self):
     def is_comp_cell(self):
         """Return *True* if the symbol is a cell in an inlined comprehension.
         """
-        return bool(self.__flags & DEF_COMP_CELL)
+        return self.is_cell() and self.__inlined_comprehension
 
     def is_namespace(self):
         """Returns *True* if name binding introduces new namespace.

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_listcomps.py

@@ -277,6 +277,20 @@ def f():
         outputs = {"y": [1]}
         self._check_in_scopes(code, outputs, scopes=["module", "function"])
 
+    def test_inlined_comp_cell_with_enclosing_free(self):
+        # The listcomp cell and the enclosing free must not share an index.
+        code = """
+            def outer(y):
+                def inner():
+                    return [lambda: x for x in (1, 2)], y
+                return inner()
+            funcs, val = outer(99)
+            z = [f() for f in funcs]
+            w = val
+        """
+        outputs = {"z": [2, 2], "w": 99}
+        self._check_in_scopes(code, outputs)
+
     def test_free_inner_cell_outer(self):
         code = """
             g = 2
@@ -407,6 +421,97 @@ def test_nested(self):
         outputs = {"y": [[0, 1], [0, 1, 4]]}
         self._check_in_scopes(code, outputs)
 
+    def test_nested_inner_uses_outer_iter(self):
+        # Inner comprehension reads the outer iteration variable. In a class
+        # this must not be treated as a class-level name of the same name.
+        code = """
+            x = 99
+            y = [[x for _ in (0,)] for x in (42,)]
+        """
+        outputs = {"y": [[42]]}
+        self._check_in_scopes(code, outputs)
+
+    def test_nested_mixed_comprehensions_use_outer_iter(self):
+        cases = [
+            ("""
+            x = 99
+            y = [{x for _ in (0,)} for x in (42,)]
+            """, {"y": [{42}]}),
+            ("""
+            x = 99
+            y = [{x: x for _ in (0,)} for x in (42,)]
+            """, {"y": [{42: 42}]}),
+            ("""
+            x = 99
+            y = {[x for _ in (0,)][0] for x in (42,)}
+            """, {"y": {42}}),
+            ("""
+            x = 99
+            y = {x: [x for _ in (0,)] for x in (42,)}
+            """, {"y": {42: [42]}}),
+        ]
+        for code, outputs in cases:
+            with self.subTest(code=code):
+                self._check_in_scopes(code, outputs)
+
+    def test_nested_triple_inner_uses_outer_iter(self):
+        code = """
+            x = 99
+            y = [[[x for _ in (0,)] for _ in (0,)] for x in (42,)]
+        """
+        outputs = {"y": [[[42]]]}
+        self._check_in_scopes(code, outputs)
+
+    def test_nested_inner_uses_outer_iter_in_iter(self):
+        code = """
+            x = 99
+            y = [[_ for _ in (x,)] for x in (42,)]
+        """
+        outputs = {"y": [[42]]}
+        self._check_in_scopes(code, outputs)
+
+    def test_nested_inner_uses_outer_iter_in_if(self):
+        code = """
+            x = 99
+            y = [[1 for _ in (0,) if x] for x in (42,)]
+        """
+        outputs = {"y": [[1]]}
+        self._check_in_scopes(code, outputs)
+
+    def test_nested_sibling_inners_use_outer_iter(self):
+        code = """
+            x = 99
+            y = [([x for _ in (0,)], [x for _ in (1,)]) for x in (42,)]
+        """
+        outputs = {"y": [([42], [42])]}
+        self._check_in_scopes(code, outputs)
+
+    def test_nested_lambda_captures_outer_iter(self):
+        code = """
+            x = 99
+            y = [[lambda: x for _ in (0,)] for x in (42,)]
+            z = y[0][0]()
+        """
+        outputs = {"z": 42}
+        self._check_in_scopes(code, outputs)
+
+    def test_nested_references___class__(self):
+        code = """
+            res = [[__class__ for _ in (0,)] for _ in (1,)]
+        """
+        self._check_in_scopes(code, raises=NameError)
+
+    def test_nested_references___class___via_lambda(self):
+        class _C:
+            res = [[lambda: __class__ for _ in (0,)] for _ in (1,)]
+        self.assertIs(_C.res[0][0](), _C)
+
+    def test_nested_references_super(self):
+        code = """
+            res = [[super for _ in (0,)] for _ in (1,)]
+        """
+        self._check_in_scopes(code, outputs={"res": [[super]]})
+
     def test_nested_2(self):
         code = """
             l = [1, 2, 3]
@@ -703,6 +808,23 @@ def test_frame_locals(self):
         """
         self._check_in_scopes(code, {"val": 0}, ns={"sys": sys})
 
+    def test_frame_locals_comp_cell_and_enclosing_free(self):
+        # The inlined listcomp cell and the enclosing free share a name.
+        # f_locals keys must still be unique so dict(**f_locals) works.
+        code = """
+            def outer(x):
+                def inner():
+                    return [(lambda: x, dict(**sys._getframe().f_locals))
+                            for x in x]
+                return inner()
+            result = outer([1, 2])
+            snaps = [d['x'] for _, d in result]
+            vals = [fn() for fn, _ in result]
+        """
+        import sys
+        self._check_in_scopes(code, {"snaps": [1, 2], "vals": [2, 2]},
+                              ns={"sys": sys}, scopes=["module", "function"])
+
     def _recursive_replace(self, maybe_code):
         if not isinstance(maybe_code, types.CodeType):
             return maybe_code