Skip to content

Commit 1c2f024

Browse files
committed
Preserve reflected multiplication for sequence subclasses
Assisted-by: Codex:GPT-6
1 parent eb2ffc6 commit 1c2f024

3 files changed

Lines changed: 133 additions & 20 deletions

File tree

‎crates/vm/src/types/slot.rs‎

Lines changed: 78 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -456,6 +456,18 @@ macro_rules! number_binary_op_wrapper {
456456
fn number_add_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
457457
vm.call_special_method(a, identifier!(vm, __add__), (b.to_owned(),))
458458
}
459+
fn number_multiply_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
460+
vm.call_special_method(a, identifier!(vm, __mul__), (b.to_owned(),))
461+
}
462+
fn number_right_multiply_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
463+
vm.call_special_method(b, identifier!(vm, __rmul__), (a.to_owned(),))
464+
}
465+
fn number_inplace_add_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
466+
vm.call_special_method(a, identifier!(vm, __iadd__), (b.to_owned(),))
467+
}
468+
fn number_inplace_multiply_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
469+
vm.call_special_method(a, identifier!(vm, __imul__), (b.to_owned(),))
470+
}
459471

460472
macro_rules! number_binary_right_op_wrapper {
461473
($name:ident) => {
@@ -1000,6 +1012,19 @@ impl PyType {
10001012
}};
10011013
}
10021014

1015+
macro_rules! native_sequence_wrapper {
1016+
($name:ident, $slot:pat) => {
1017+
self.get_attr(identifier!(ctx, $name)).is_some_and(|attr| {
1018+
attr.class().is(ctx.types.wrapper_descriptor_type)
1019+
&& attr.downcast_ref::<PyWrapper>().is_some_and(|descr| {
1020+
descr.name.as_str() == stringify!($name)
1021+
&& self.mro.read().iter().any(|cls| cls.is(descr.typ))
1022+
&& matches!(descr.wrapped, $slot)
1023+
})
1024+
})
1025+
};
1026+
}
1027+
10031028
match accessor {
10041029
// === Main slots ===
10051030
SlotAccessor::TpRepr => update_main_slot!(repr, repr_wrapper, Repr),
@@ -1291,12 +1316,22 @@ impl PyType {
12911316
}
12921317
}
12931318
SlotAccessor::NbInplaceAdd => {
1294-
update_sub_slot!(
1295-
as_number,
1296-
inplace_add,
1297-
number_binary_op_wrapper!(__iadd__),
1298-
NumBinary
1299-
)
1319+
if native_sequence_wrapper!(__iadd__, SlotFunc::SeqConcat(_)) {
1320+
// Unlike repeat wrappers, both inplace-add slots accept an
1321+
// object argument. Keep the numeric dispatcher ahead of
1322+
// reflected addition, including after deleting an override.
1323+
self.slots
1324+
.as_number
1325+
.inplace_add
1326+
.store(Some(number_inplace_add_wrapper));
1327+
} else {
1328+
update_sub_slot!(
1329+
as_number,
1330+
inplace_add,
1331+
number_inplace_add_wrapper,
1332+
NumBinary
1333+
)
1334+
}
13001335
}
13011336
SlotAccessor::NbSubtract => {
13021337
if name.as_str() == "__rsub__" {
@@ -1324,30 +1359,55 @@ impl PyType {
13241359
)
13251360
}
13261361
SlotAccessor::NbMultiply => {
1327-
if name.as_str() == "__rmul__" {
1362+
let native_left = native_sequence_wrapper!(__mul__, SlotFunc::SeqRepeat(_));
1363+
let native_right = native_sequence_wrapper!(__rmul__, SlotFunc::SeqRepeat(_));
1364+
// CPython has a single numeric multiply slot for both names.
1365+
// Keep native repeat in sq_repeat unless either name needs the
1366+
// numeric dispatcher, or that dispatcher was already installed.
1367+
let needs_numeric = self.slots.as_number.multiply.load().is_some()
1368+
|| self.slots.as_number.right_multiply.load().is_some()
1369+
|| self.slots.as_sequence.repeat.load().is_none()
1370+
|| (!native_left && self.get_attr(identifier!(ctx, __mul__)).is_some())
1371+
|| (!native_right && self.get_attr(identifier!(ctx, __rmul__)).is_some());
1372+
if native_right {
1373+
self.slots
1374+
.as_number
1375+
.right_multiply
1376+
.store(needs_numeric.then_some(number_right_multiply_wrapper));
1377+
} else if name.as_str() == "__rmul__" {
13281378
update_sub_slot!(
13291379
as_number,
13301380
right_multiply,
1331-
number_binary_right_op_wrapper!(__rmul__),
1381+
number_right_multiply_wrapper,
13321382
NumBinary
13331383
)
1384+
}
1385+
if native_left {
1386+
self.slots
1387+
.as_number
1388+
.multiply
1389+
.store(needs_numeric.then_some(number_multiply_wrapper));
1390+
} else if name.as_str() == "__mul__" {
1391+
update_sub_slot!(as_number, multiply, number_multiply_wrapper, NumBinary)
1392+
}
1393+
}
1394+
SlotAccessor::NbInplaceMultiply => {
1395+
if native_sequence_wrapper!(__imul__, SlotFunc::SeqRepeat(_)) {
1396+
let needs_numeric = self.slots.as_number.inplace_multiply.load().is_some()
1397+
|| self.slots.as_sequence.inplace_repeat.load().is_none();
1398+
self.slots
1399+
.as_number
1400+
.inplace_multiply
1401+
.store(needs_numeric.then_some(number_inplace_multiply_wrapper));
13341402
} else {
13351403
update_sub_slot!(
13361404
as_number,
1337-
multiply,
1338-
number_binary_op_wrapper!(__mul__),
1405+
inplace_multiply,
1406+
number_inplace_multiply_wrapper,
13391407
NumBinary
13401408
)
13411409
}
13421410
}
1343-
SlotAccessor::NbInplaceMultiply => {
1344-
update_sub_slot!(
1345-
as_number,
1346-
inplace_multiply,
1347-
number_binary_op_wrapper!(__imul__),
1348-
NumBinary
1349-
)
1350-
}
13511411
SlotAccessor::NbRemainder => {
13521412
if name.as_str() == "__rmod__" {
13531413
update_sub_slot!(

‎crates/vm/src/types/slot_defs.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -653,7 +653,11 @@ impl SlotAccessor {
653653
// Number slots
654654
Self::NbAdd => copy_number!(add),
655655
Self::NbSubtract => copy_number!(subtract),
656-
Self::NbMultiply => copy_number!(multiply),
656+
Self::NbMultiply => {
657+
// CPython inherits one nb_multiply slot for both directions.
658+
copy_number!(multiply);
659+
copy_number!(right_multiply);
660+
}
657661
Self::NbRemainder => copy_number!(remainder),
658662
Self::NbDivmod => copy_number!(divmod),
659663
Self::NbPower => copy_number!(power),

‎extra_tests/snippets/operator_arithmetic.py‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from operator import imul, mul
1+
from operator import iadd, imul, mul
22

33
from testutils import assert_raises
44

@@ -253,3 +253,52 @@ class RepeatInt(int):
253253
assert_raises(TypeError, operation, value)
254254
delattr(RepeatList, method)
255255
assert operation(value) == [1, 1]
256+
257+
258+
# Inherited sequence slots must leave reflected numeric dispatch its priority.
259+
class ReflectedCount(int):
260+
def __rmul__(self, value):
261+
return "reflected"
262+
263+
264+
class InheritedRepeat(list):
265+
pass
266+
267+
268+
for operation in (mul, imul):
269+
value = InheritedRepeat([1])
270+
assert operation(value, ReflectedCount(2)) == "reflected"
271+
assert value == [1]
272+
273+
274+
class RepeatMixin:
275+
def __rmul__(self, value):
276+
return "mixin"
277+
278+
279+
class MixedRepeat(list, RepeatMixin):
280+
pass
281+
282+
283+
# A secondary base can supply the numeric slot even if list hides its method.
284+
assert mul(MixedRepeat([1]), ReflectedCount(2)) == [1, 1]
285+
assert imul(ReflectedCount(2), MixedRepeat([1])) == [1, 1]
286+
287+
# Removing an override preserves an already installed numeric dispatcher.
288+
InheritedRepeat.__imul__ = lambda self, other: NotImplemented
289+
del InheritedRepeat.__imul__
290+
value = InheritedRepeat([1])
291+
assert imul(value, ReflectedCount(2)) is value
292+
assert value == [1, 1]
293+
294+
295+
class ReflectedAdd(list):
296+
def __radd__(self, value):
297+
return "reflected"
298+
299+
300+
InheritedRepeat.__iadd__ = lambda self, other: NotImplemented
301+
del InheritedRepeat.__iadd__
302+
value = InheritedRepeat([1])
303+
assert iadd(value, ReflectedAdd([2])) is value
304+
assert value == [1, 2]

0 commit comments

Comments
 (0)