From e02aae5f45bc04dbd58ddee47c0b3ef933fa92d2 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 15 Oct 2026 09:31:06 +0100 Subject: [PATCH 1/2] [mypyc] Evaluate args in correct order when builtin takes a genexpr When a generator expression is passed to sum(), min() or another builtin that mypyc has special handling for, match CPython's order of evaluation. --- mypyc/irbuild/for_helpers.py | 70 +++++++++++++++++-- mypyc/irbuild/specialize.py | 55 ++++++++------- mypyc/test-data/fixtures/ir.py | 14 +++- mypyc/test-data/irbuild-basic.test | 18 ++--- mypyc/test-data/run-generator-exprs.test | 89 ++++++++++++++++++++++++ mypyc/test/test_run.py | 1 + 6 files changed, 208 insertions(+), 39 deletions(-) create mode 100644 mypyc/test-data/run-generator-exprs.test diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 74db465183379..072736f5fda31 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -106,6 +106,7 @@ def for_loop_helper( else_insts: GenFunc | None, is_async: bool, line: int, + after_init: GenFunc | None = None, ) -> None: """Generate IR for a loop. @@ -114,6 +115,8 @@ def for_loop_helper( expr: the expression to iterate over body_insts: a function that generates the body of the loop else_insts: a function that generates the else block instructions + after_init: a function that generates instructions to run after the iterable has + been evaluated (and iter() called on it) but before the first iteration """ # Body of the loop body_block = BasicBlock() @@ -130,6 +133,8 @@ def for_loop_helper( for_gen = make_for_loop_generator( builder, index, expr, body_block, normal_loop_exit, line, is_async=is_async ) + if after_init is not None: + after_init() builder.push_loop_stack(step_block, exit_block) condition_block = BasicBlock() @@ -223,6 +228,7 @@ def sequence_from_generator_preallocate_helper( gen: GeneratorExpr, empty_op_llbuilder: Callable[[Value, int], Value], set_item_op: Callable[[Value, Value, Value, int], None], + after_first_init: GenFunc | None = None, ) -> Value | None: """Generate a new tuple or list from a simple generator expression. @@ -245,6 +251,8 @@ def sequence_from_generator_preallocate_helper( - New Value See `new_list_set_item_op` and `new_tuple_set_item_op` for detailed implementation. + after_first_init: Optional function that generates instructions to run after the + sequence has been evaluated (see comprehension_helper). """ if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0: line = gen.line @@ -289,6 +297,9 @@ def sequence_from_generator_preallocate_helper( else: sequence = builder.accept(sequence_expr) + if after_first_init is not None: + after_first_init() + length = get_expr_length_value(builder, sequence_expr, sequence, line, use_pyssize_t=True) target_op = empty_op_llbuilder(length, line) @@ -306,8 +317,16 @@ def set_item(item_index: Value) -> None: return None -def translate_list_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value: +def translate_list_comprehension( + builder: IRBuilder, gen: GeneratorExpr, *, after_first_init: GenFunc | None = None +) -> Value: + """Generate a list from a comprehension or generator expression. + + See comprehension_helper for the meaning of after_first_init. + """ if raise_error_if_contains_unreachable_names(builder, gen): + if after_first_init is not None: + after_first_init() return builder.none() def set_item(x: Value, y: Value, z: Value, line: int) -> None: @@ -319,6 +338,7 @@ def set_item(x: Value, y: Value, z: Value, line: int) -> None: gen, empty_op_llbuilder=builder.builder.new_list_op_with_length, set_item_op=set_item, + after_first_init=after_first_init, ) if val is not None: return val @@ -331,10 +351,45 @@ def gen_inner_stmts() -> None: e = builder.accept(gen.left_expr) builder.primitive_op(list_append_op, [builder.read(list_ops, gen.line), e], gen.line) - comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line) + comprehension_helper( + builder, loop_params, gen_inner_stmts, gen.line, after_first_init=after_first_init + ) return builder.read(list_ops, gen.line) +def translate_generator_to_list( + builder: IRBuilder, gen: GeneratorExpr, *, before_loop: GenFunc | None = None +) -> Value: + """Generate a list from a generator expression whose items are all consumed at once. + + This is equivalent to creating and fully consuming a generator object, when the + consumer doesn't run other code while iterating. + + If given, before_loop generates code for the other arguments of the consuming call. + CPython evaluates them after the first iterable (when the generator object is created) + but before the generator runs, so we do the same. + """ + return translate_list_comprehension(builder, gen, after_first_init=before_loop) + + +def generator_loop_helper( + builder: IRBuilder, + gen: GeneratorExpr, + gen_inner_stmts: GenFunc, + *, + before_loop: GenFunc | None = None, +) -> None: + """Generate loops that iterate a generator expression inline (no generator object). + + gen_inner_stmts generates code for each item (gen.left_expr). See + translate_generator_to_list for the meaning of before_loop. + """ + loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async)) + comprehension_helper( + builder, loop_params, gen_inner_stmts, gen.line, after_first_init=before_loop + ) + + def raise_error_if_contains_unreachable_names( builder: IRBuilder, gen: GeneratorExpr | DictionaryComprehension ) -> bool: @@ -407,6 +462,7 @@ def comprehension_helper( loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]], gen_inner_stmts: Callable[[], None], line: int, + after_first_init: GenFunc | None = None, ) -> None: """Helper function for list comprehensions. @@ -417,9 +473,14 @@ def comprehension_helper( - "conditions" is a list of conditions, evaluated in order with short-circuiting, that must all be true for the loop body to be executed gen_inner_stmts: function to generate the IR for the body of the innermost loop + after_first_init: function to generate IR to run after the first (outermost) + iterable has been evaluated, but before the first iteration """ - def handle_loop(loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]]) -> None: + def handle_loop( + loop_params: list[tuple[Lvalue, Expression, list[Expression], bool]], + after_init: GenFunc | None = None, + ) -> None: """Generate IR for a loop. Given a list of (index, expression, [conditions]) tuples, generate IR @@ -434,6 +495,7 @@ def handle_loop(loop_params: list[tuple[Lvalue, Expression, list[Expression], bo None, is_async=is_async, line=line, + after_init=after_init, ) def loop_contents( @@ -468,7 +530,7 @@ def loop_contents( with builder.enter_borrow_scope(line): gen_inner_stmts() - handle_loop(loop_params) + handle_loop(loop_params, after_first_init) def is_range_ref(expr: RefExpr) -> bool: diff --git a/mypyc/irbuild/specialize.py b/mypyc/irbuild/specialize.py index 637c3bca75c38..27f461acb7859 100644 --- a/mypyc/irbuild/specialize.py +++ b/mypyc/irbuild/specialize.py @@ -94,9 +94,10 @@ from mypyc.irbuild.constant_fold import constant_fold_expr from mypyc.irbuild.for_helpers import ( comprehension_helper, + generator_loop_helper, get_expr_length_value, sequence_from_generator_preallocate_helper, - translate_list_comprehension, + translate_generator_to_list, translate_set_comprehension, ) from mypyc.irbuild.format_str_tokenizer import ( @@ -561,28 +562,29 @@ def translate_safe_generator_call( and expr.arg_kinds[0] == ARG_POS and isinstance(expr.args[0], GeneratorExpr) ): - if isinstance(callee, MemberExpr): + receiver = builder.accept(callee.expr) if isinstance(callee, MemberExpr) else None + other_args: list[Value] = [] + + def gen_other_args() -> None: + other_args.extend(builder.accept(arg) for arg in expr.args[1:]) + + items = translate_generator_to_list( + builder, expr.args[0], before_loop=gen_other_args if len(expr.args) > 1 else None + ) + args = [items] + other_args + if receiver is not None: + assert isinstance(callee, MemberExpr) return builder.gen_method_call( - builder.accept(callee.expr), + receiver, callee.name, - ( - [translate_list_comprehension(builder, expr.args[0])] - + [builder.accept(arg) for arg in expr.args[1:]] - ), + args, builder.node_type(expr), expr.line, expr.arg_kinds, expr.arg_names, ) else: - return builder.call_refexpr_with_args( - expr, - callee, - ( - [translate_list_comprehension(builder, expr.args[0])] - + [builder.accept(arg) for arg in expr.args[1:]] - ), - ) + return builder.call_refexpr_with_args(expr, callee, args) return None @@ -667,9 +669,10 @@ def translate_sum_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> V gen_expr = expr.args[0] target_type = builder.node_type(expr) retval = Register(target_type, line=expr.line) - builder.assign( - retval, builder.coerce(builder.accept(start_expr), target_type, expr.line), expr.line - ) + + def gen_start() -> None: + start = builder.coerce(builder.accept(start_expr), target_type, expr.line) + builder.assign(retval, start, expr.line) def gen_inner_stmts() -> None: call_expr = builder.accept(gen_expr.left_expr) @@ -677,10 +680,7 @@ def gen_inner_stmts() -> None: retval, builder.binary_op(retval, call_expr, "+", call_expr.line), call_expr.line ) - loop_params = list( - zip(gen_expr.indices, gen_expr.sequences, gen_expr.condlists, gen_expr.is_async) - ) - comprehension_helper(builder, loop_params, gen_inner_stmts, gen_expr.line) + generator_loop_helper(builder, gen_expr, gen_inner_stmts, before_loop=gen_start) return retval @@ -720,17 +720,22 @@ def translate_next_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> gen = expr.args[0] retval = Register(builder.node_type(expr)) - default_val = builder.accept(expr.args[1]) if len(expr.args) > 1 else None + default_val: Value | None = None exit_block = BasicBlock() + def gen_default() -> None: + nonlocal default_val + default_val = builder.accept(expr.args[1]) + def gen_inner_stmts() -> None: # next takes the first element of the generator, so if # something gets produced, we are done. builder.assign(retval, builder.accept(gen.left_expr), gen.left_expr.line) builder.goto(exit_block) - loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async)) - comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line) + generator_loop_helper( + builder, gen, gen_inner_stmts, before_loop=gen_default if len(expr.args) > 1 else None + ) # Now we need the case for when nothing got hit. If there was # a default value, we produce it, and otherwise we raise diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py index 3cc3feb24e0e7..8fed798bbc695 100644 --- a/mypyc/test-data/fixtures/ir.py +++ b/mypyc/test-data/fixtures/ir.py @@ -442,10 +442,22 @@ def pow(base: __SupportsPow2[T_contra, T_co], exp: T_contra, mod: None = None) - def pow(base: __SupportsPow3NoneOnly[T_contra, T_co], exp: T_contra, mod: None = None) -> T_co: ... @overload def pow(base: __SupportsPow3[T_contra, _M, T_co], exp: T_contra, mod: _M) -> T_co: ... -def sorted(iterable: Iterable[_T]) -> list[_T]: ... +def sorted( + iterable: Iterable[_T], *, key: Callable[[_T], Any] | None = None, reverse: bool = False +) -> list[_T]: ... def exit() -> None: ... +@overload def min(x: _T, y: _T) -> _T: ... +@overload +def min( + iterable: Iterable[_T], *, key: Callable[[_T], Any] | None = None, default: _T = ... +) -> _T: ... +@overload def max(x: _T, y: _T) -> _T: ... +@overload +def max( + iterable: Iterable[_T], *, key: Callable[[_T], Any] | None = None, default: _T = ... +) -> _T: ... def repr(o: object) -> str: ... def ascii(o: object) -> str: ... def ord(o: object) -> int: ... diff --git a/mypyc/test-data/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test index f22f9afff3d32..83b2a16be7914 100644 --- a/mypyc/test-data/irbuild-basic.test +++ b/mypyc/test-data/irbuild-basic.test @@ -2944,9 +2944,9 @@ def call_sum(l: Iterable[int], comparison: Callable[[int], bool]) -> int: [out] def call_sum(l, comparison): - l, comparison :: object - r0 :: int - r1, r2 :: object + l, comparison, r0 :: object + r1 :: int + r2 :: object r3, x :: int r4 :: object r5 :: object[1] @@ -2956,10 +2956,10 @@ def call_sum(l, comparison): r10, r11 :: int r12 :: bit L0: - r0 = 0 - r1 = PyObject_GetIter(l) + r0 = PyObject_GetIter(l) + r1 = 0 L1: - r2 = PyIter_Next(r1) + r2 = PyIter_Next(r0) if is_error(r2) goto L4 else goto L2 L2: r3 = unbox(int, r2) @@ -2972,14 +2972,14 @@ L2: r8 = unbox(bool, r7) r9 = r8 << 1 r10 = extend r9: builtins.bool to builtins.int - r11 = CPyTagged_Add(r0, r10) - r0 = r11 + r11 = CPyTagged_Add(r1, r10) + r1 = r11 L3: goto L1 L4: r12 = CPy_NoErrOccurred() L5: - return r0 + return r1 [case testSetAttr1] from typing import Any, Dict, List diff --git a/mypyc/test-data/run-generator-exprs.test b/mypyc/test-data/run-generator-exprs.test new file mode 100644 index 0000000000000..6fcb1cef2442e --- /dev/null +++ b/mypyc/test-data/run-generator-exprs.test @@ -0,0 +1,89 @@ +-- Generator expressions + +[case testGeneratorExprConsumerArgumentOrder] +from collections import OrderedDict +from typing import Any + +log: list[str] = [] + +def it() -> list[int]: + log.append("iterable") + return [1, 2] + +def elt(x: int) -> int: + log.append(f"elt {x}") + return x + +def arg(name: str, value: Any) -> Any: + log.append(name) + return value + +def check(expected: list[str]) -> None: + actual = log[:] + log.clear() + assert actual == expected, actual + +def test_sum() -> None: + assert sum((elt(x) for x in it()), arg("start", 10)) == 13 + check(["iterable", "start", "elt 1", "elt 2"]) + assert sum((elt(x) for x in it() if x > 1), start=arg("start", 10)) == 12 + check(["iterable", "start", "elt 2"]) + assert sum((elt(x) + y for x in it() for y in it()), arg("start", 0)) == 12 + check(["iterable", "start", "iterable", "elt 1", "elt 1", "iterable", "elt 2", "elt 2"]) + +def test_sum_over_range() -> None: + n = 3 + assert sum((elt(x) for x in range(n)), arg("start", 1)) == 4 + check(["start", "elt 0", "elt 1", "elt 2"]) + +def test_next() -> None: + assert next((elt(x) for x in it()), arg("default", -1)) == 1 + check(["iterable", "default", "elt 1"]) + assert next((elt(x) for x in it() if x > 5), arg("default", -1)) == -1 + check(["iterable", "default"]) + +def test_min_max_sorted() -> None: + assert max((elt(x) for x in it()), key=arg("key", abs)) == 2 + check(["iterable", "key", "elt 1", "elt 2"]) + assert min((elt(x) for x in it() if x > 5), default=arg("default", 0)) == 0 + check(["iterable", "default"]) + assert sorted((elt(x) for x in it()), key=arg("key", lambda x: -x)) == [2, 1] + check(["iterable", "key", "elt 1", "elt 2"]) + +def test_dict() -> None: + expected = {"1": 1, "2": 2, "k": 0} + assert dict(((str(elt(x)), x) for x in it()), k=arg("kwarg", 0)) == expected + check(["iterable", "kwarg", "elt 1", "elt 2"]) + assert OrderedDict(((str(elt(x)), x) for x in it()), k=arg("kwarg", 0)) == expected + check(["iterable", "kwarg", "elt 1", "elt 2"]) + +def receiver(d: dict[str, int]) -> dict[str, int]: + log.append("receiver") + return d + +def test_method() -> None: + d: dict[str, int] = {} + receiver(d).update(((str(elt(x)), x) for x in it()), k=arg("kwarg", 0)) + assert d == {"1": 1, "2": 2, "k": 0} + check(["receiver", "iterable", "kwarg", "elt 1", "elt 2"]) + +def test_exception_in_later_argument() -> None: + def fail() -> Any: + log.append("fail") + raise ValueError("fail") + + try: + sum((elt(x) for x in it()), fail()) + except ValueError: + pass + check(["iterable", "fail"]) + try: + next((elt(x) for x in it()), fail()) + except ValueError: + pass + check(["iterable", "fail"]) + try: + max((elt(x) for x in it()), key=fail()) + except ValueError: + pass + check(["iterable", "fail"]) diff --git a/mypyc/test/test_run.py b/mypyc/test/test_run.py index 868f83654b47c..55923dfae3aa2 100644 --- a/mypyc/test/test_run.py +++ b/mypyc/test/test_run.py @@ -66,6 +66,7 @@ "run-classes.test", "run-traits.test", "run-generators.test", + "run-generator-exprs.test", "run-generics.test", "run-multimodule.test", "run-bench.test", From 5d5f1711f281eb89c90d8d4501231b999db0e170 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Fri, 16 Oct 2026 15:19:29 +0100 Subject: [PATCH 2/2] Address feedback --- mypyc/test-data/run-generator-exprs.test | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/mypyc/test-data/run-generator-exprs.test b/mypyc/test-data/run-generator-exprs.test index 6fcb1cef2442e..17d0f97e164b6 100644 --- a/mypyc/test-data/run-generator-exprs.test +++ b/mypyc/test-data/run-generator-exprs.test @@ -4,6 +4,8 @@ from collections import OrderedDict from typing import Any +from testutil import assertRaises + log: list[str] = [] def it() -> list[int]: @@ -72,18 +74,12 @@ def test_exception_in_later_argument() -> None: log.append("fail") raise ValueError("fail") - try: + with assertRaises(ValueError, "fail"): sum((elt(x) for x in it()), fail()) - except ValueError: - pass check(["iterable", "fail"]) - try: + with assertRaises(ValueError, "fail"): next((elt(x) for x in it()), fail()) - except ValueError: - pass check(["iterable", "fail"]) - try: + with assertRaises(ValueError, "fail"): max((elt(x) for x in it()), key=fail()) - except ValueError: - pass check(["iterable", "fail"])