diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index fc81320fa91..3cf9fb7fd6f 100644 --- 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 index 6445c655306..3408c47daf7 100644 --- 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()