From 684e800396099183f83379c41880b37b499a02bc Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Thu, 25 Sep 2025 20:54:04 -0400 Subject: [PATCH 01/19] [mypyc] feat: extend get_expr_length for enumerate, map, zip, and range This PR is pretty simple, I just extended get_expr_length to work for a few more obvious cases: - `builtins.enumerate` - `builtins.map` - `builtins.zip` - `builtins.range` This PR is ready for review. Are you going to want tests for all of these? I don't want to spend time now until I know for sure. --- mypyc/irbuild/for_helpers.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 5edee6cb4df4..348765e690dd 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1200,10 +1200,23 @@ def get_expr_length(expr: Expression) -> int | None: and expr.node.has_explicit_value ): return len(expr.node.final_value) + elif isinstance(expr, CallExpr) and isinstance(callee := expr.callee, NameExpr): + fullname = callee.fullname + if fullname == "builtins.enumerate" and len(expr.args) == 1: + return get_expr_length(expr.args[0]) + elif fullname == "builtins.map" and len(expr.args) == 2: + return get_expr_length(expr.args[1]) + elif fullname == "builtins.zip" and expr.args: + arg_lengths = [get_expr_length(arg) for arg in expr.args] + if all(arg is not None for arg in arg_lengths): + return min(arg_lengths) + elif fullname == "builtins.range" and all(isinstance(arg, IntExpr) for arg in expr.args): + return len(range(*(arg.value for arg in expr.args))) + # TODO: extend this, passing length of listcomp and genexp should have worthwhile # performance boost and can be (sometimes) figured out pretty easily. set and dict # comps *can* be done as well but will need special logic to consider the possibility - # of key conflicts. Range, enumerate, zip are all simple logic. + # of key conflicts. return None From c792337d2f98166dae6986385e4c835d1980f39d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Sep 2025 00:55:59 +0000 Subject: [PATCH 02/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 348765e690dd..2dd00d21861e 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1212,7 +1212,7 @@ def get_expr_length(expr: Expression) -> int | None: return min(arg_lengths) elif fullname == "builtins.range" and all(isinstance(arg, IntExpr) for arg in expr.args): return len(range(*(arg.value for arg in expr.args))) - + # TODO: extend this, passing length of listcomp and genexp should have worthwhile # performance boost and can be (sometimes) figured out pretty easily. set and dict # comps *can* be done as well but will need special logic to consider the possibility From 63196179910b15ece7dac4c3e97b6ec742cb8aa2 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Thu, 25 Sep 2025 21:04:06 -0400 Subject: [PATCH 03/19] fix: missing import --- mypyc/irbuild/for_helpers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 2dd00d21861e..dc516ec72ab8 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -16,6 +16,7 @@ DictionaryComprehension, Expression, GeneratorExpr, + IntExpr, ListExpr, Lvalue, MemberExpr, From 48ebc0dc3b408e2c7d9f4dc86d75be716a60ceff Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Thu, 25 Sep 2025 21:12:18 -0400 Subject: [PATCH 04/19] fix mypy errs --- mypyc/irbuild/for_helpers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index dc516ec72ab8..00de814d77a5 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1210,9 +1210,9 @@ def get_expr_length(expr: Expression) -> int | None: elif fullname == "builtins.zip" and expr.args: arg_lengths = [get_expr_length(arg) for arg in expr.args] if all(arg is not None for arg in arg_lengths): - return min(arg_lengths) + return min(arg_lengths) # type: ignore [type-var] elif fullname == "builtins.range" and all(isinstance(arg, IntExpr) for arg in expr.args): - return len(range(*(arg.value for arg in expr.args))) + return len(range(*(arg.value for arg in expr.args))) # type: ignore [attr-defined] # TODO: extend this, passing length of listcomp and genexp should have worthwhile # performance boost and can be (sometimes) figured out pretty easily. set and dict From 3e31189752435041d5190195c00e5714172e92ef Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Thu, 25 Sep 2025 22:07:00 -0400 Subject: [PATCH 05/19] support list and tuple call expr --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 00de814d77a5..f9e724006b39 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1203,7 +1203,7 @@ def get_expr_length(expr: Expression) -> int | None: return len(expr.node.final_value) elif isinstance(expr, CallExpr) and isinstance(callee := expr.callee, NameExpr): fullname = callee.fullname - if fullname == "builtins.enumerate" and len(expr.args) == 1: + if fullname in ("builtins.list", "builtins.tuple", "builtins.enumerate") and len(expr.args) == 1: return get_expr_length(expr.args[0]) elif fullname == "builtins.map" and len(expr.args) == 2: return get_expr_length(expr.args[1]) From d38d4b4d5becb0190cfa2beb7e469778932caed8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Sep 2025 02:08:21 +0000 Subject: [PATCH 06/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index f9e724006b39..186d545ae08b 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1203,7 +1203,10 @@ def get_expr_length(expr: Expression) -> int | None: return len(expr.node.final_value) elif isinstance(expr, CallExpr) and isinstance(callee := expr.callee, NameExpr): fullname = callee.fullname - if fullname in ("builtins.list", "builtins.tuple", "builtins.enumerate") and len(expr.args) == 1: + if ( + fullname in ("builtins.list", "builtins.tuple", "builtins.enumerate") + and len(expr.args) == 1 + ): return get_expr_length(expr.args[0]) elif fullname == "builtins.map" and len(expr.args) == 2: return get_expr_length(expr.args[1]) From 4940cfc5f805b4c525479378a0a47a8cad444f4e Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Thu, 25 Sep 2025 23:19:28 -0400 Subject: [PATCH 07/19] validate arg kind --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 186d545ae08b..d83ddc10100c 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1201,7 +1201,7 @@ def get_expr_length(expr: Expression) -> int | None: and expr.node.has_explicit_value ): return len(expr.node.final_value) - elif isinstance(expr, CallExpr) and isinstance(callee := expr.callee, NameExpr): + elif isinstance(expr, CallExpr) and isinstance(callee := expr.callee, NameExpr) and all(kind == ARG_POS for kind in expr.arg_kinds): fullname = callee.fullname if ( fullname in ("builtins.list", "builtins.tuple", "builtins.enumerate") From 50179797a631a006cdca443865a2d8325f7a9405 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Sep 2025 03:20:47 +0000 Subject: [PATCH 08/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index d83ddc10100c..7ebc0b26684b 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1201,7 +1201,11 @@ def get_expr_length(expr: Expression) -> int | None: and expr.node.has_explicit_value ): return len(expr.node.final_value) - elif isinstance(expr, CallExpr) and isinstance(callee := expr.callee, NameExpr) and all(kind == ARG_POS for kind in expr.arg_kinds): + elif ( + isinstance(expr, CallExpr) + and isinstance(callee := expr.callee, NameExpr) + and all(kind == ARG_POS for kind in expr.arg_kinds) + ): fullname = callee.fullname if ( fullname in ("builtins.list", "builtins.tuple", "builtins.enumerate") From b6f098dfa5367540b05c5406143891d6d9a559b9 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Thu, 25 Sep 2025 23:21:59 -0400 Subject: [PATCH 09/19] support sorted and reversed --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 7ebc0b26684b..ddcd5e0ca267 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1208,7 +1208,7 @@ def get_expr_length(expr: Expression) -> int | None: ): fullname = callee.fullname if ( - fullname in ("builtins.list", "builtins.tuple", "builtins.enumerate") + fullname in ("builtins.list", "builtins.tuple", "builtins.enumerate", "builtins.sorted", "builtins.reversed") and len(expr.args) == 1 ): return get_expr_length(expr.args[0]) From 3389573a39bb64fd4cf2a115c8bd52f6907ce626 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Sep 2025 03:23:25 +0000 Subject: [PATCH 10/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index ddcd5e0ca267..cceece120cc7 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1208,7 +1208,14 @@ def get_expr_length(expr: Expression) -> int | None: ): fullname = callee.fullname if ( - fullname in ("builtins.list", "builtins.tuple", "builtins.enumerate", "builtins.sorted", "builtins.reversed") + fullname + in ( + "builtins.list", + "builtins.tuple", + "builtins.enumerate", + "builtins.sorted", + "builtins.reversed", + ) and len(expr.args) == 1 ): return get_expr_length(expr.args[0]) From e788168bfed81a6b8b75045eb3751d855b63fcc8 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 26 Sep 2025 07:37:17 -0400 Subject: [PATCH 11/19] testTupleBuiltFromLengthCheckable --- mypyc/test-data/irbuild-tuple.test | 47 ++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/mypyc/test-data/irbuild-tuple.test b/mypyc/test-data/irbuild-tuple.test index 081cc1b174c9..b4333202ed82 100644 --- a/mypyc/test-data/irbuild-tuple.test +++ b/mypyc/test-data/irbuild-tuple.test @@ -829,6 +829,53 @@ L4: a = r1 return 1 +[case testTupleBuiltFromLengthCheckable] +from typing import Tuple + +def f(val: bool) -> bool: + return not val + +def test() -> None: + a = tuple((i, x) for i, (_, x) in zip(map(str, range(5)), enumerate(sorted(reversed("abc"))))) +[out] +def f(val): + val, r0 :: bool +L0: + r0 = val ^ 1 + return r0 +def test(source): + source :: tuple + r0 :: native_int + r1 :: tuple + r2 :: native_int + r3 :: bit + r4 :: object + r5, x, r6 :: bool + r7 :: object + r8 :: native_int + a :: tuple +L0: + r0 = var_object_size source + r1 = PyTuple_New(r0) + r2 = 0 +L1: + r3 = r2 < r0 :: signed + if r3 goto L2 else goto L4 :: bool +L2: + r4 = CPySequenceTuple_GetItemUnsafe(source, r2) + r5 = unbox(bool, r4) + x = r5 + r6 = f(x) + r7 = box(bool, r6) + CPySequenceTuple_SetItemUnsafe(r1, r2, r7) +L3: + r8 = r2 + 1 + r2 = r8 + goto L1 +L4: + a = r1 + return 1 + [case testTupleBuiltFromStars] from typing import Final From 68c81ce2e68093fdbb3fb192f1ca04b0d16bb2f3 Mon Sep 17 00:00:00 2001 From: BobTheBuidler Date: Fri, 26 Sep 2025 12:54:53 +0000 Subject: [PATCH 12/19] fix: add map fixture --- mypyc/test-data/fixtures/ir.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py index a4b4f3ce2b1f..55529ba51772 100644 --- a/mypyc/test-data/fixtures/ir.py +++ b/mypyc/test-data/fixtures/ir.py @@ -308,6 +308,11 @@ def __iter__(self) -> Iterator[int]: pass def __len__(self) -> int: pass def __next__(self) -> int: pass +class map(Iterator[_S]): + def __init__(self, func: Callable[[_T], _S], iterable: Iterable[_T]) -> None: pass + def __iter__(self) -> Self: pass + def __next__(self) -> _S: pass + class property: def __init__(self, fget: Optional[Callable[[Any], Any]] = ..., fset: Optional[Callable[[Any, Any], None]] = ..., From 5f4ef4df7b80eedfa542f43e0d1003768af4a368 Mon Sep 17 00:00:00 2001 From: BobTheBuidler Date: Fri, 26 Sep 2025 12:57:59 +0000 Subject: [PATCH 13/19] fix: IR --- mypyc/test-data/irbuild-tuple.test | 125 +++++++++++++++++++++++------ 1 file changed, 99 insertions(+), 26 deletions(-) diff --git a/mypyc/test-data/irbuild-tuple.test b/mypyc/test-data/irbuild-tuple.test index b4333202ed82..4e9422d9c603 100644 --- a/mypyc/test-data/irbuild-tuple.test +++ b/mypyc/test-data/irbuild-tuple.test @@ -836,44 +836,117 @@ def f(val: bool) -> bool: return not val def test() -> None: - a = tuple((i, x) for i, (_, x) in zip(map(str, range(5)), enumerate(sorted(reversed("abc"))))) + # this tuple is created from a very complex genexp but we can still compute the length and preallocate the tuple + a = tuple( + x + for x + in zip( + map(str, range(5)), + enumerate(sorted(reversed(tuple("abcdefg")))) + ) + ) [out] def f(val): val, r0 :: bool L0: r0 = val ^ 1 return r0 -def test(source): - source :: tuple - r0 :: native_int - r1 :: tuple - r2 :: native_int - r3 :: bit - r4 :: object - r5, x, r6 :: bool - r7 :: object - r8 :: native_int - a :: tuple +def test(): + r0 :: list + r1, r2, r3 :: object + r4 :: object[1] + r5 :: object_ptr + r6 :: object + r7 :: range + r8 :: object + r9 :: str + r10 :: object + r11 :: object[2] + r12 :: object_ptr + r13 :: object + r14 :: str + r15 :: tuple + r16 :: object + r17 :: str + r18 :: object + r19 :: object[1] + r20 :: object_ptr + r21 :: object + r22 :: list + r23 :: object + r24 :: str + r25 :: object + r26 :: object[1] + r27 :: object_ptr + r28, r29 :: object + r30 :: str + r31 :: object + r32 :: object[2] + r33 :: object_ptr + r34, r35, r36 :: object + r37, x :: tuple[str, tuple[int, str]] + r38 :: object + r39 :: i32 + r40, r41 :: bit + r42, a :: tuple L0: - r0 = var_object_size source - r1 = PyTuple_New(r0) - r2 = 0 + r0 = PyList_New(0) + r1 = load_address PyUnicode_Type + r2 = load_address PyRange_Type + r3 = object 5 + r4 = [r3] + r5 = load_address r4 + r6 = PyObject_Vectorcall(r2, r5, 1, 0) + keep_alive r3 + r7 = cast(range, r6) + r8 = builtins :: module + r9 = 'map' + r10 = CPyObject_GetAttr(r8, r9) + r11 = [r1, r7] + r12 = load_address r11 + r13 = PyObject_Vectorcall(r10, r12, 2, 0) + keep_alive r1, r7 + r14 = 'abcdefg' + r15 = PySequence_Tuple(r14) + r16 = builtins :: module + r17 = 'reversed' + r18 = CPyObject_GetAttr(r16, r17) + r19 = [r15] + r20 = load_address r19 + r21 = PyObject_Vectorcall(r18, r20, 1, 0) + keep_alive r15 + r22 = CPySequence_Sort(r21) + r23 = builtins :: module + r24 = 'enumerate' + r25 = CPyObject_GetAttr(r23, r24) + r26 = [r22] + r27 = load_address r26 + r28 = PyObject_Vectorcall(r25, r27, 1, 0) + keep_alive r22 + r29 = builtins :: module + r30 = 'zip' + r31 = CPyObject_GetAttr(r29, r30) + r32 = [r13, r28] + r33 = load_address r32 + r34 = PyObject_Vectorcall(r31, r33, 2, 0) + keep_alive r13, r28 + r35 = PyObject_GetIter(r34) L1: - r3 = r2 < r0 :: signed - if r3 goto L2 else goto L4 :: bool + r36 = PyIter_Next(r35) + if is_error(r36) goto L4 else goto L2 L2: - r4 = CPySequenceTuple_GetItemUnsafe(source, r2) - r5 = unbox(bool, r4) - x = r5 - r6 = f(x) - r7 = box(bool, r6) - CPySequenceTuple_SetItemUnsafe(r1, r2, r7) + r37 = unbox(tuple[str, tuple[int, str]], r36) + x = r37 + r38 = box(tuple[str, tuple[int, str]], x) + r39 = PyList_Append(r0, r38) + r40 = r39 >= 0 :: signed L3: - r8 = r2 + 1 - r2 = r8 goto L1 L4: - a = r1 + r41 = CPy_NoErrOccurred() +L5: + r42 = PyList_AsTuple(r0) + a = r42 return 1 [case testTupleBuiltFromStars] From a64f2e4ef2e1121f4376cf114ad50f65548ce568 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Fri, 26 Sep 2025 17:28:34 -0400 Subject: [PATCH 14/19] fix: check arg count for range --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index cceece120cc7..a0c6bfd8579d 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1225,7 +1225,7 @@ def get_expr_length(expr: Expression) -> int | None: arg_lengths = [get_expr_length(arg) for arg in expr.args] if all(arg is not None for arg in arg_lengths): return min(arg_lengths) # type: ignore [type-var] - elif fullname == "builtins.range" and all(isinstance(arg, IntExpr) for arg in expr.args): + elif fullname == "builtins.range" and len(expr.args) <= 3 and all(isinstance(arg, IntExpr) for arg in expr.args): return len(range(*(arg.value for arg in expr.args))) # type: ignore [attr-defined] # TODO: extend this, passing length of listcomp and genexp should have worthwhile From f6f98e42020c783de664058ca345477e0609a81c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Sep 2025 21:30:00 +0000 Subject: [PATCH 15/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index a0c6bfd8579d..cd07688593bb 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1225,7 +1225,11 @@ def get_expr_length(expr: Expression) -> int | None: arg_lengths = [get_expr_length(arg) for arg in expr.args] if all(arg is not None for arg in arg_lengths): return min(arg_lengths) # type: ignore [type-var] - elif fullname == "builtins.range" and len(expr.args) <= 3 and all(isinstance(arg, IntExpr) for arg in expr.args): + elif ( + fullname == "builtins.range" + and len(expr.args) <= 3 + and all(isinstance(arg, IntExpr) for arg in expr.args) + ): return len(range(*(arg.value for arg in expr.args))) # type: ignore [attr-defined] # TODO: extend this, passing length of listcomp and genexp should have worthwhile From e376ec3dda2af673c4cd738c110456fa15c484f0 Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Sat, 4 Oct 2025 01:31:10 -0400 Subject: [PATCH 16/19] constant fold range args --- mypyc/irbuild/for_helpers.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index ba9303d11594..16de50e4b1e0 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -7,6 +7,7 @@ from __future__ import annotations +import contextlib from typing import Callable, ClassVar from mypy.nodes import ( @@ -68,6 +69,7 @@ short_int_rprimitive, ) from mypyc.irbuild.builder import IRBuilder +from mypyc.irbuild.constant_fold import constant_fold_expr from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple from mypyc.primitives.dict_ops import ( @@ -1204,18 +1206,18 @@ def gen_cleanup(self) -> None: gen.gen_cleanup() -def get_expr_length(expr: Expression) -> int | None: +def get_expr_length(builder: IRBuilder, expr: Expression) -> int | None: if isinstance(expr, (StrExpr, BytesExpr)): return len(expr.value) elif isinstance(expr, (ListExpr, TupleExpr)): # if there are no star expressions, or we know the length of them, # we know the length of the expression - stars = [get_expr_length(i) for i in expr.items if isinstance(i, StarExpr)] + stars = [get_expr_length(builder, i) for i in expr.items if isinstance(i, StarExpr)] if None not in stars: other = sum(not isinstance(i, StarExpr) for i in expr.items) return other + sum(stars) # type: ignore [arg-type] elif isinstance(expr, StarExpr): - return get_expr_length(expr.expr) + return get_expr_length(builder, expr.expr) elif ( isinstance(expr, RefExpr) and isinstance(expr.node, Var) @@ -1241,19 +1243,20 @@ def get_expr_length(expr: Expression) -> int | None: ) and len(expr.args) == 1 ): - return get_expr_length(expr.args[0]) + return get_expr_length(builder, expr.args[0]) elif fullname == "builtins.map" and len(expr.args) == 2: - return get_expr_length(expr.args[1]) + return get_expr_length(builder, expr.args[1]) elif fullname == "builtins.zip" and expr.args: - arg_lengths = [get_expr_length(arg) for arg in expr.args] + arg_lengths = [get_expr_length(builder, arg) for arg in expr.args] if all(arg is not None for arg in arg_lengths): return min(arg_lengths) # type: ignore [type-var] - elif ( - fullname == "builtins.range" - and len(expr.args) <= 3 - and all(isinstance(arg, IntExpr) for arg in expr.args) - ): - return len(range(*(arg.value for arg in expr.args))) # type: ignore [attr-defined] + elif fullname == "builtins.range" and len(expr.args) <= 3: + folded_args = [constant_fold_expr(builder, arg) for arg in args] + if all(isinstance(arg, int) for arg in folded_args): + try: + return len(range(*folded_args)) + except ValueError: # prevent crash if invalid args + pass # TODO: extend this, passing length of listcomp and genexp should have worthwhile # performance boost and can be (sometimes) figured out pretty easily. set and dict From c398aa73664fc77ae2ea49df7015f8a75e0aed82 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Oct 2025 05:38:25 +0000 Subject: [PATCH 17/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mypyc/irbuild/for_helpers.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 16de50e4b1e0..093f06390211 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -7,7 +7,6 @@ from __future__ import annotations -import contextlib from typing import Callable, ClassVar from mypy.nodes import ( @@ -17,7 +16,6 @@ DictionaryComprehension, Expression, GeneratorExpr, - IntExpr, ListExpr, Lvalue, MemberExpr, From 293f00117d8c9138fab4e34624e8c7c23c54e56d Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Sat, 4 Oct 2025 01:55:45 -0400 Subject: [PATCH 18/19] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 093f06390211..ae9c96abc88e 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1268,7 +1268,7 @@ def get_expr_length_value( ) -> Value: rtype = builder.node_type(expr) assert is_sequence_rprimitive(rtype) or isinstance(rtype, RTuple), rtype - length = get_expr_length(expr) + length = get_expr_length(builder, expr) if length is None: # We cannot compute the length at compile time, so we will fetch it. return builder.builder.builtin_len(expr_reg, line, use_pyssize_t=use_pyssize_t) From d0a6e7c75d39719574fb5455fd32f980b2580ffb Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Mon, 6 Oct 2025 13:38:29 -0400 Subject: [PATCH 19/19] Update for_helpers.py --- mypyc/irbuild/for_helpers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index ae9c96abc88e..8e708178d1df 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import Callable, ClassVar +from typing import Callable, ClassVar, cast from mypy.nodes import ( ARG_POS, @@ -1249,10 +1249,10 @@ def get_expr_length(builder: IRBuilder, expr: Expression) -> int | None: if all(arg is not None for arg in arg_lengths): return min(arg_lengths) # type: ignore [type-var] elif fullname == "builtins.range" and len(expr.args) <= 3: - folded_args = [constant_fold_expr(builder, arg) for arg in args] + folded_args = [constant_fold_expr(builder, arg) for arg in expr.args] if all(isinstance(arg, int) for arg in folded_args): try: - return len(range(*folded_args)) + return len(range(*cast(list[int], folded_args))) except ValueError: # prevent crash if invalid args pass