pythongh-97933: inline sync list/dict/set comprehensions

Author: carljm

Doc/library/dis.rst

@@ -1214,6 +1214,16 @@ iterations of the loop.
 
    .. versionadded:: 3.12
 
+.. opcode:: LOAD_FAST_OR_NULL (var_num)
+
+   Pushes a reference to the local ``co_varnames[var_num]`` onto the stack, or
+   pushes ``NULL`` onto the stack if the local variable has not been
+   initialized. This opcode has the same runtime effect as ``LOAD_FAST``; it
+   exists to maintain the invariant that ``LOAD_FAST`` will never load ``NULL``
+   and may appear only where the variable is guaranteed to be initialized.
+
+   .. versionadded:: 3.12
+
 .. opcode:: STORE_FAST (var_num)
 
    Stores ``STACK.pop()`` into the local ``co_varnames[var_num]``.

Include/internal/pycore_opcode.h

@@ -64,6 +64,7 @@ typedef struct _symtable_entry {
     unsigned ste_needs_class_closure : 1; /* for class scopes, true if a
                                              closure over __class__
                                              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 */
     int ste_comp_iter_expr; /* non-zero if visiting a comprehension range expression */
     int ste_lineno;          /* first line of block */

Include/internal/pycore_symtable.h

@@ -431,6 +431,7 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.12a5 3515 (Embed jump mask in COMPARE_OP oparg)
 #     Python 3.12a5 3516 (Add COMPARE_AND_BRANCH instruction)
 #     Python 3.12a5 3517 (Change YIELD_VALUE oparg to exception block depth)
+#     Python 3.12a5 3518 (Inline sync list/dict/set comprehensions in the calling function)
 
 #     Python 3.13 will start with 3550
 
@@ -443,7 +444,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3517).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3518).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Include/opcode.h

@@ -196,6 +196,8 @@ def pseudo_op(name, op, real_ops):
 hascompare.append(141)
 
 def_op('CALL_FUNCTION_EX', 142)  # Flags
+def_op('LOAD_FAST_OR_NULL', 143)  # Local variable number, may load NULL if undefined
+haslocal.append(143)
 
 def_op('EXTENDED_ARG', 144)
 EXTENDED_ARG = 144

Lib/importlib/_bootstrap_external.py

@@ -161,15 +161,15 @@ def bug708901():
 
 
 def bug1333982(x=[]):
-    assert 0, ([s for s in x] +
+    assert 0, ((s for s in x) +
               1)
     pass
 
 dis_bug1333982 = """\
 %3d        RESUME                   0
 
 %3d        LOAD_ASSERTION_ERROR
-           LOAD_CONST               1 (<code object <listcomp> at 0x..., file "%s", line %d>)
+           LOAD_CONST               1 (<code object <genexpr> at 0x..., file "%s", line %d>)
            MAKE_FUNCTION            0
            LOAD_FAST                0 (x)
            GET_ITER
@@ -642,7 +642,7 @@ async def _co(x):
 def _h(y):
     def foo(x):
         '''funcdoc'''
-        return [x + z for z in y]
+        return list(x + z for z in y)
     return foo
 
 dis_nested_0 = """\
@@ -672,13 +672,15 @@ def foo(x):
 
 %3d        RESUME                   0
 
-%3d        LOAD_CLOSURE             0 (x)
+%3d        LOAD_GLOBAL              1 (NULL + list)
+           LOAD_CLOSURE             0 (x)
            BUILD_TUPLE              1
-           LOAD_CONST               1 (<code object <listcomp> at 0x..., file "%s", line %d>)
+           LOAD_CONST               1 (<code object <genexpr> at 0x..., file "%s", line %d>)
            MAKE_FUNCTION            8 (closure)
            LOAD_DEREF               1 (y)
            GET_ITER
            CALL                     0
+           CALL                     1
            RETURN_VALUE
 """ % (dis_nested_0,
        __file__,
@@ -690,21 +692,29 @@ def foo(x):
 )
 
 dis_nested_2 = """%s
-Disassembly of <code object <listcomp> at 0x..., file "%s", line %d>:
+Disassembly of <code object <genexpr> at 0x..., file "%s", line %d>:
            COPY_FREE_VARS           1
 
-%3d        RESUME                   0
-           BUILD_LIST               0
+%3d        RETURN_GENERATOR
+           POP_TOP
+           RESUME                   0
            LOAD_FAST                0 (.0)
-        >> FOR_ITER                 7 (to 26)
+        >> FOR_ITER                 9 (to 32)
            STORE_FAST               1 (z)
            LOAD_DEREF               2 (x)
            LOAD_FAST                1 (z)
            BINARY_OP                0 (+)
-           LIST_APPEND              2
-           JUMP_BACKWARD            9 (to 8)
+           YIELD_VALUE              1
+           RESUME                   1
+           POP_TOP
+           JUMP_BACKWARD           11 (to 10)
         >> END_FOR
+           LOAD_CONST               0 (None)
            RETURN_VALUE
+        >> CALL_INTRINSIC_1         3
+           RERAISE                  1
+ExceptionTable:
+1 row
 """ % (dis_nested_1,
        __file__,
        _h.__code__.co_firstlineno + 3,

Lib/opcode.py

@@ -4106,14 +4106,14 @@ def test(*args, **kwargs):
 
     @cpython_only
     def test_signature_bind_implicit_arg(self):
-        # Issue #19611: getcallargs should work with set comprehensions
+        # Issue #19611: getcallargs should work with comprehensions
         def make_set():
-            return {z * z for z in range(5)}
-        setcomp_code = make_set.__code__.co_consts[1]
-        setcomp_func = types.FunctionType(setcomp_code, {})
+            return set(z * z for z in range(5))
+        gencomp_code = make_set.__code__.co_consts[1]
+        gencomp_func = types.FunctionType(gencomp_code, {})
 
         iterator = iter(range(5))
-        self.assertEqual(self.call(setcomp_func, iterator), {0, 1, 4, 9, 16})
+        self.assertEqual(set(self.call(gencomp_func, iterator)), {0, 1, 4, 9, 16})
 
     def test_signature_bind_posonly_kwargs(self):
         def foo(bar, /, **kwargs):

Lib/test/test_dis.py

@@ -143,6 +143,48 @@
     >>> test_func()
     [2, 2, 2, 2, 2]
 
+A comprehension's iteration var, if in a cell, doesn't stomp on a previous value:
+
+    >>> def test_func():
+    ...     y = 10
+    ...     items = [(lambda: y) for y in range(5)]
+    ...     x = y
+    ...     y = 20
+    ...     return x, [z() for z in items]
+    >>> test_func()
+    (10, [4, 4, 4, 4, 4])
+
+A comprehension's iteration var doesn't shadow implicit globals for sibling scopes:
+
+    >>> g = -1
+    >>> def test_func():
+    ...     def inner():
+    ...         return g
+    ...     [g for g in range(5)]
+    ...     return inner
+    >>> test_func()()
+    -1
+
+A modification to a closed-over variable is visible in the outer scope:
+
+    >>> def test_func():
+    ...     x = -1
+    ...     items = [(x:=y) for y in range(3)]
+    ...     return x
+    >>> test_func()
+    2
+
+Comprehensions' scopes don't interact with each other:
+
+    >>> def test_func():
+    ...     lst = list(range(3))
+    ...     ret = [lambda: x for x in lst]
+    ...     inc = [x + 1 for x in lst]
+    ...     [x for x in inc]
+    ...     return ret
+    >>> test_func()[0]()
+    2
+
 """
 
 

Lib/test/test_inspect.py

@@ -187,9 +187,7 @@ def test_trace_list_comprehension(self):
         firstlineno_called = get_firstlineno(traced_doubler)
         expected = {
             (self.my_py_filename, firstlineno_calling + 1): 1,
-            # List comprehensions work differently in 3.x, so the count
-            # below changed compared to 2.x.
-            (self.my_py_filename, firstlineno_calling + 2): 12,
+            (self.my_py_filename, firstlineno_calling + 2): 11,
             (self.my_py_filename, firstlineno_calling + 3): 1,
             (self.my_py_filename, firstlineno_called + 1): 10,
         }