diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index f1536fb3111a..30101371bb3f 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,25 @@ 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 + 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: builder = self.builder line = self.line diff --git a/mypyc/primitives/generic_ops.py b/mypyc/primitives/generic_ops.py index 87b1f4ff309b..738f7bf8808a 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 8c5c95c8223e..7aed23a3b530 100644 --- a/mypyc/test-data/irbuild-statements.test +++ b/mypyc/test-data/irbuild-statements.test @@ -78,6 +78,73 @@ 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 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 243d029a137b..077acafdb814 100644 --- a/mypyc/test-data/run-loops.test +++ b/mypyc/test-data/run-loops.test @@ -572,6 +572,91 @@ 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_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)] + +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] + 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_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] + 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