Fix sub_table ordering for nested inlined comprehensions (PEP 709) (#7480)

LegNeato · web-flow · commit 2180f535d850 · 2026-03-22T17:23:03.000+09:00
When an inlined comprehension's first iterator expression contains
nested scopes (such as a lambda), those scopes' sub_tables appear at the
current position in the parent's sub_table list. The previous code
spliced the comprehension's own child sub_tables (e.g. inner inlined
comprehensions) into that same position before compiling the iterator,
which shifted the iterator's sub_tables to wrong indices.

Move the splice after the first iterator is compiled so its sub_tables
are consumed at their original positions.

Fixes nested list comprehensions like:
```python
    [[x for _, x in g] for _, g in itertools.groupby(..., lambda x: ...)]
```

Disclosure: I used AI to develop the patch though I was heavily
involved.
diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs
@@ -8013,23 +8013,31 @@ impl Compiler {
     ) -> CompileResult<()> {
         // PEP 709: Consume the comprehension's sub_table.
         // The symbols are already merged into parent scope by analyze_symbol_table.
-        // Splice the comprehension's children into the parent so nested scopes
-        // (e.g. inner comprehensions, lambdas) can be found by the compiler.
         let current_table = self
             .symbol_table_stack
             .last_mut()
             .expect("no current symbol table");
         let comp_table = current_table.sub_tables[current_table.next_sub_table].clone();
         current_table.next_sub_table += 1;
+
+        // Compile the outermost iterator first. Its expression may reference
+        // nested scopes (e.g. lambdas) whose sub_tables sit at the current
+        // position in the parent's list. Those must be consumed before we
+        // splice in the comprehension's own children.
+        self.compile_expression(&generators[0].iter)?;
+
+        // Splice the comprehension's children (e.g. nested inlined
+        // comprehensions) into the parent so the compiler can find them.
         if !comp_table.sub_tables.is_empty() {
+            let current_table = self
+                .symbol_table_stack
+                .last_mut()
+                .expect("no current symbol table");
             let insert_pos = current_table.next_sub_table;
             for (i, st) in comp_table.sub_tables.iter().enumerate() {
                 current_table.sub_tables.insert(insert_pos + i, st.clone());
             }
         }
-
-        // Step 1: Compile the outermost iterator BEFORE tweaking scopes
-        self.compile_expression(&generators[0].iter)?;
         if has_async && generators[0].is_async {
             emit!(self, Instruction::GetAIter);
         } else {
diff --git a/extra_tests/snippets/syntax_comprehension.py b/extra_tests/snippets/syntax_comprehension.py
@@ -45,3 +45,37 @@
 def f():
     # Test no panic occurred.
     [[x := 1 for j in range(5)] for i in range(5)]
+
+
+# Nested inlined comprehensions with lambda in the first iterator expression.
+# The lambda's sub_table must be consumed before the inner comprehension's
+# sub_table is spliced in, otherwise scope ordering is wrong.
+def test_nested_comp_with_lambda():
+    import itertools
+    offsets = {0: [0], 1: [1], 3: [2]}
+    grouped = [
+        [x for _, x in group]
+        for _, group in itertools.groupby(
+            enumerate(sorted(offsets.keys())), lambda x: x[1] - x[0]
+        )
+    ]
+    assert grouped == [[0, 1], [3]], f"got {grouped}"
+
+test_nested_comp_with_lambda()
+
+
+# Nested inlined comprehensions with throwaway `_` in both levels.
+def test_nested_comp_underscore():
+    data = [(1, "a", "x"), (2, "b", "y")]
+    result = [[v for _, v in zip(range(2), row)] for _, *row in data]
+    assert result == [["a", "x"], ["b", "y"]], f"got {result}"
+
+test_nested_comp_underscore()
+
+
+# Simple nested inlined comprehensions.
+def test_simple_nested_comp():
+    result = [[j * i for j in range(3)] for i in range(3)]
+    assert result == [[0, 0, 0], [0, 1, 2], [0, 2, 4]]
+
+test_simple_nested_comp()
