From 3bee9a8df1deb756e6340ae2b23077be38d3bf5d Mon Sep 17 00:00:00 2001 From: Ryan Heard Date: Sun, 4 Oct 2026 03:51:13 -0400 Subject: [PATCH 1/3] [mypyc] Use __index__ for range() arguments in for loops --- mypyc/irbuild/for_helpers.py | 17 ++++++++- mypyc/primitives/generic_ops.py | 8 ++++ mypyc/test-data/irbuild-statements.test | 33 ++++++++++++++++ mypyc/test-data/run-loops.test | 50 +++++++++++++++++++++++++ 4 files changed, 107 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index f1536fb3111a6..7c8af58f3e0bf 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -54,6 +54,7 @@ bool_rprimitive, c_pyssize_t_rprimitive, int_rprimitive, + is_bool_or_bit_rprimitive, is_dict_rprimitive, is_fixed_width_rtype, is_immutable_rprimitive, @@ -61,6 +62,7 @@ is_sequence_rprimitive, is_short_int_rprimitive, is_str_rprimitive, + is_tagged, is_tuple_rprimitive, object_pointer_rprimitive, object_rprimitive, @@ -80,7 +82,7 @@ dict_value_iter_op, ) from mypyc.primitives.exc_ops import no_err_occurred_op, propagate_if_error_op -from mypyc.primitives.generic_ops import aiter_op, anext_op, iter_op, next_op +from mypyc.primitives.generic_ops import aiter_op, anext_op, index_op, iter_op, next_op from mypyc.primitives.list_ops import ( list_append_op, list_get_item_int64_op, @@ -1122,6 +1124,8 @@ class ForRange(ForGenerator): def init(self, start_reg: Value, end_reg: Value, step: int) -> None: builder = self.builder + start_reg = self.convert_arg(start_reg) + end_reg = self.convert_arg(end_reg) self.start_reg = start_reg self.end_reg = end_reg self.step = step @@ -1139,6 +1143,17 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None: self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index) builder.assign(self.index_target, builder.read(self.index_reg, self.line), self.line) + def convert_arg(self, value: Value) -> Value: + """Convert a range() argument to an int using __index__, like range() does.""" + if ( + is_tagged(value.type) + or is_fixed_width_rtype(value.type) + or is_bool_or_bit_rprimitive(value.type) + ): + return value + index = self.builder.call_c(index_op, [value], self.line) + return self.builder.coerce(index, int_rprimitive, self.line) + def gen_condition(self) -> None: builder = self.builder line = self.line diff --git a/mypyc/primitives/generic_ops.py b/mypyc/primitives/generic_ops.py index 87b1f4ff309b2..738f7bf8808ac 100644 --- a/mypyc/primitives/generic_ops.py +++ b/mypyc/primitives/generic_ops.py @@ -178,6 +178,14 @@ priority=0, ) +# operator.index(obj) +index_op = custom_op( + arg_types=[object_rprimitive], + return_type=object_rprimitive, + c_function_name="PyNumber_Index", + error_kind=ERR_MAGIC, +) + # obj1[obj2] py_get_item_op = method_op( name="__getitem__", diff --git a/mypyc/test-data/irbuild-statements.test b/mypyc/test-data/irbuild-statements.test index 8c5c95c8223ec..5c1bc1c0e3dec 100644 --- a/mypyc/test-data/irbuild-statements.test +++ b/mypyc/test-data/irbuild-statements.test @@ -78,6 +78,39 @@ L3: L4: return 1 +[case testForInRangeAnyArgs] +from typing import Any + +def f(a: Any, b: Any) -> None: + for i in range(a, b): + pass +[out] +def f(a, b): + a, b, r0 :: object + r1 :: int + r2 :: object + r3, r4, i :: int + r5 :: bit + r6 :: int +L0: + r0 = PyNumber_Index(a) + r1 = unbox(int, r0) + r2 = PyNumber_Index(b) + r3 = unbox(int, r2) + r4 = r1 + i = r4 +L1: + r5 = int_lt r4, r3 + if r5 goto L2 else goto L4 :: bool +L2: + i = r4 +L3: + r6 = CPyTagged_Add(r4, 2) + r4 = r6 + goto L1 +L4: + return 1 + [case testBreak] def f() -> None: n = 0 diff --git a/mypyc/test-data/run-loops.test b/mypyc/test-data/run-loops.test index 243d029a137b0..c11e0e7956fd0 100644 --- a/mypyc/test-data/run-loops.test +++ b/mypyc/test-data/run-loops.test @@ -572,6 +572,56 @@ def test_range_object() -> None: except ValueError as e: assert "range() arg 3 must not be zero" in str(e) +[case testRangeLoopIndexArgs] +from typing import Any + +from testutil import assertRaises + +def loop_end(end: Any) -> list[int]: + result = [] + for i in range(end): + result.append(i) + return result + +def loop(start: Any, end: Any) -> list[int]: + return [i for i in range(start, end)] + +def loop_down(start: Any) -> list[int]: + return [i for i in range(start, 0, -1)] + +def loop_bool(b: bool) -> list[int]: + return [i for i in range(b)] + +def test_index_args() -> None: + from index_types import Index + assert loop_end(Index(3)) == [0, 1, 2] + assert loop(Index(1), Index(3)) == [1, 2] + assert loop_down(Index(2)) == [2, 1] + +def test_int_args() -> None: + assert loop_end(3) == [0, 1, 2] + assert loop(True, 3) == [1, 2] + n = 2**100 + assert loop(n, n + 2) == [n, n + 1] + assert loop_bool(True) == [0] + assert loop_bool(False) == [] + +def test_invalid_args() -> None: + with assertRaises(TypeError, "'float' object cannot be interpreted as an integer"): + loop_end(2.5) + with assertRaises(TypeError, "'str' object cannot be interpreted as an integer"): + loop(0, "3") + with assertRaises(TypeError, "'float' object cannot be interpreted as an integer"): + loop_down(2.0) + +[file index_types.py] +class Index: + def __init__(self, n: int) -> None: + self.n = n + + def __index__(self) -> int: + return self.n + [case testNamedTupleLoop] from collections.abc import Iterable from typing import NamedTuple, Any From 3416aefd84b06133970c59916022e4f8557737c9 Mon Sep 17 00:00:00 2001 From: Ryan Heard Date: Wed, 7 Oct 2026 12:14:23 -0400 Subject: [PATCH 2/3] [mypyc] Add test for big integers --- mypyc/test-data/run-loops.test | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/mypyc/test-data/run-loops.test b/mypyc/test-data/run-loops.test index c11e0e7956fd0..6eb88441fd06d 100644 --- a/mypyc/test-data/run-loops.test +++ b/mypyc/test-data/run-loops.test @@ -589,6 +589,9 @@ def loop(start: Any, end: Any) -> list[int]: def loop_down(start: Any) -> list[int]: return [i for i in range(start, 0, -1)] +def loop_down_to(start: Any, end: Any) -> list[int]: + return [i for i in range(start, end, -1)] + def loop_bool(b: bool) -> list[int]: return [i for i in range(b)] @@ -598,6 +601,17 @@ def test_index_args() -> None: assert loop(Index(1), Index(3)) == [1, 2] assert loop_down(Index(2)) == [2, 1] +def test_big_index_args() -> None: + from index_types import Index + # Ints >= 2**62 or < -2**62 don't fit in a short int, so these loops cross + # between short and heap-allocated values + n = 2**62 + assert loop(Index(n - 1), Index(n + 2)) == [n - 1, n, n + 1] + assert loop(Index(-n - 2), Index(-n + 1)) == [-n - 2, -n - 1, -n] + assert loop_down_to(Index(n + 1), Index(n - 2)) == [n + 1, n, n - 1] + assert loop_down_to(Index(-n + 1), Index(-n - 2)) == [-n + 1, -n, -n - 1] + assert loop_end(Index(-n - 1)) == [] + def test_int_args() -> None: assert loop_end(3) == [0, 1, 2] assert loop(True, 3) == [1, 2] From 46d34ac7c298892cea4eef1cc06c6f7358496aad Mon Sep 17 00:00:00 2001 From: Ryan Heard Date: Wed, 7 Oct 2026 12:17:30 -0400 Subject: [PATCH 3/3] [mypyc] Directly call the __index__ method on native classes that have it --- mypyc/irbuild/for_helpers.py | 10 +++++++- mypyc/test-data/irbuild-statements.test | 34 +++++++++++++++++++++++++ mypyc/test-data/run-loops.test | 21 +++++++++++++++ 3 files changed, 64 insertions(+), 1 deletion(-) diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 7c8af58f3e0bf..30101371bb3f7 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -1151,7 +1151,15 @@ def convert_arg(self, value: Value) -> Value: or is_bool_or_bit_rprimitive(value.type) ): return value - index = self.builder.call_c(index_op, [value], self.line) + if ( + isinstance(value.type, RInstance) + and value.type.class_ir.is_ext_class + and value.type.class_ir.has_method("__index__") + ): + # Directly call the __index__ method on native classes that have it. + index = self.builder.gen_method_call(value, "__index__", [], int_rprimitive, self.line) + else: + index = self.builder.call_c(index_op, [value], self.line) return self.builder.coerce(index, int_rprimitive, self.line) def gen_condition(self) -> None: diff --git a/mypyc/test-data/irbuild-statements.test b/mypyc/test-data/irbuild-statements.test index 5c1bc1c0e3dec..7aed23a3b5307 100644 --- a/mypyc/test-data/irbuild-statements.test +++ b/mypyc/test-data/irbuild-statements.test @@ -111,6 +111,40 @@ L3: L4: return 1 +[case testForInRangeNativeIndexArg] +class C: + def __index__(self) -> int: + return 2 + +def f(c: C) -> None: + for i in range(c): # type: ignore[arg-type] + pass +[out] +def C.__index__(self): + self :: __main__.C +L0: + return 4 +def f(c): + c :: __main__.C + r0, r1, i :: int + r2 :: bit + r3 :: int +L0: + r0 = c.__index__() + r1 = 0 + i = r1 +L1: + r2 = int_lt r1, r0 + if r2 goto L2 else goto L4 :: bool +L2: + i = r1 +L3: + r3 = CPyTagged_Add(r1, 2) + r1 = r3 + goto L1 +L4: + return 1 + [case testBreak] def f() -> None: n = 0 diff --git a/mypyc/test-data/run-loops.test b/mypyc/test-data/run-loops.test index 6eb88441fd06d..077acafdb8144 100644 --- a/mypyc/test-data/run-loops.test +++ b/mypyc/test-data/run-loops.test @@ -595,6 +595,21 @@ def loop_down_to(start: Any, end: Any) -> list[int]: def loop_bool(b: bool) -> list[int]: return [i for i in range(b)] +class NativeIndex: + def __init__(self, n: int) -> None: + self.n = n + + def __index__(self) -> int: + return self.n + +class NativeIndexSub(NativeIndex): + def __index__(self) -> int: + return self.n + 1 + +def loop_native(start: NativeIndex, end: NativeIndex) -> list[int]: + # The test fixtures only accept int arguments for range() + return [i for i in range(start, end)] # type: ignore[arg-type] + def test_index_args() -> None: from index_types import Index assert loop_end(Index(3)) == [0, 1, 2] @@ -612,6 +627,12 @@ def test_big_index_args() -> None: assert loop_down_to(Index(-n + 1), Index(-n - 2)) == [-n + 1, -n, -n - 1] assert loop_end(Index(-n - 1)) == [] +def test_native_index_args() -> None: + assert loop_native(NativeIndex(1), NativeIndex(3)) == [1, 2] + assert loop_native(NativeIndexSub(1), NativeIndex(4)) == [2, 3] + n = 2**62 + assert loop_native(NativeIndex(n - 1), NativeIndex(n + 1)) == [n - 1, n] + def test_int_args() -> None: assert loop_end(3) == [0, 1, 2] assert loop(True, 3) == [1, 2]