Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 78 additions & 18 deletions crates/vm/src/types/slot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -456,6 +456,18 @@ macro_rules! number_binary_op_wrapper {
fn number_add_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
vm.call_special_method(a, identifier!(vm, __add__), (b.to_owned(),))
}
fn number_multiply_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
vm.call_special_method(a, identifier!(vm, __mul__), (b.to_owned(),))
}
fn number_right_multiply_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
vm.call_special_method(b, identifier!(vm, __rmul__), (a.to_owned(),))
}
fn number_inplace_add_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
vm.call_special_method(a, identifier!(vm, __iadd__), (b.to_owned(),))
}
fn number_inplace_multiply_wrapper(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
vm.call_special_method(a, identifier!(vm, __imul__), (b.to_owned(),))
}

macro_rules! number_binary_right_op_wrapper {
($name:ident) => {
Expand Down Expand Up @@ -1000,6 +1012,19 @@ impl PyType {
}};
}

macro_rules! native_sequence_wrapper {
($name:ident, $slot:pat) => {
self.get_attr(identifier!(ctx, $name)).is_some_and(|attr| {
attr.class().is(ctx.types.wrapper_descriptor_type)
&& attr.downcast_ref::<PyWrapper>().is_some_and(|descr| {
descr.name.as_str() == stringify!($name)
&& self.mro.read().iter().any(|cls| cls.is(descr.typ))
&& matches!(descr.wrapped, $slot)
})
})
};
}

match accessor {
// === Main slots ===
SlotAccessor::TpRepr => update_main_slot!(repr, repr_wrapper, Repr),
Expand Down Expand Up @@ -1291,12 +1316,22 @@ impl PyType {
}
}
SlotAccessor::NbInplaceAdd => {
update_sub_slot!(
as_number,
inplace_add,
number_binary_op_wrapper!(__iadd__),
NumBinary
)
if native_sequence_wrapper!(__iadd__, SlotFunc::SeqConcat(_)) {
// Unlike repeat wrappers, both inplace-add slots accept an
// object argument. Keep the numeric dispatcher ahead of
// reflected addition, including after deleting an override.
self.slots
.as_number
.inplace_add
.store(Some(number_inplace_add_wrapper));
} else {
update_sub_slot!(
as_number,
inplace_add,
number_inplace_add_wrapper,
NumBinary
)
}
}
SlotAccessor::NbSubtract => {
if name.as_str() == "__rsub__" {
Expand Down Expand Up @@ -1324,30 +1359,55 @@ impl PyType {
)
}
SlotAccessor::NbMultiply => {
if name.as_str() == "__rmul__" {
let native_left = native_sequence_wrapper!(__mul__, SlotFunc::SeqRepeat(_));
let native_right = native_sequence_wrapper!(__rmul__, SlotFunc::SeqRepeat(_));
// CPython has a single numeric multiply slot for both names.
// Keep native repeat in sq_repeat unless either name needs the
// numeric dispatcher, or that dispatcher was already installed.
let needs_numeric = self.slots.as_number.multiply.load().is_some()
|| self.slots.as_number.right_multiply.load().is_some()
|| self.slots.as_sequence.repeat.load().is_none()
|| (!native_left && self.get_attr(identifier!(ctx, __mul__)).is_some())
|| (!native_right && self.get_attr(identifier!(ctx, __rmul__)).is_some());
if native_right {
self.slots
.as_number
.right_multiply
.store(needs_numeric.then_some(number_right_multiply_wrapper));
} else if name.as_str() == "__rmul__" {
update_sub_slot!(
as_number,
right_multiply,
number_binary_right_op_wrapper!(__rmul__),
number_right_multiply_wrapper,
NumBinary
)
}
if native_left {
self.slots
.as_number
.multiply
.store(needs_numeric.then_some(number_multiply_wrapper));
} else if name.as_str() == "__mul__" {
update_sub_slot!(as_number, multiply, number_multiply_wrapper, NumBinary)
}
}
SlotAccessor::NbInplaceMultiply => {
if native_sequence_wrapper!(__imul__, SlotFunc::SeqRepeat(_)) {
let needs_numeric = self.slots.as_number.inplace_multiply.load().is_some()
|| self.slots.as_sequence.inplace_repeat.load().is_none();
self.slots
.as_number
.inplace_multiply
.store(needs_numeric.then_some(number_inplace_multiply_wrapper));
} else {
update_sub_slot!(
as_number,
multiply,
number_binary_op_wrapper!(__mul__),
inplace_multiply,
number_inplace_multiply_wrapper,
NumBinary
)
}
}
SlotAccessor::NbInplaceMultiply => {
update_sub_slot!(
as_number,
inplace_multiply,
number_binary_op_wrapper!(__imul__),
NumBinary
)
}
SlotAccessor::NbRemainder => {
if name.as_str() == "__rmod__" {
update_sub_slot!(
Expand Down
6 changes: 5 additions & 1 deletion crates/vm/src/types/slot_defs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -653,7 +653,11 @@ impl SlotAccessor {
// Number slots
Self::NbAdd => copy_number!(add),
Self::NbSubtract => copy_number!(subtract),
Self::NbMultiply => copy_number!(multiply),
Self::NbMultiply => {
// CPython inherits one nb_multiply slot for both directions.
copy_number!(multiply);
copy_number!(right_multiply);
}
Self::NbRemainder => copy_number!(remainder),
Self::NbDivmod => copy_number!(divmod),
Self::NbPower => copy_number!(power),
Expand Down
51 changes: 50 additions & 1 deletion extra_tests/snippets/operator_arithmetic.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from operator import imul, mul
from operator import iadd, imul, mul

from testutils import assert_raises

Expand Down Expand Up @@ -253,3 +253,52 @@ class RepeatInt(int):
assert_raises(TypeError, operation, value)
delattr(RepeatList, method)
assert operation(value) == [1, 1]


# Inherited sequence slots must leave reflected numeric dispatch its priority.
class ReflectedCount(int):
def __rmul__(self, value):
return "reflected"


class InheritedRepeat(list):
pass


for operation in (mul, imul):
value = InheritedRepeat([1])
assert operation(value, ReflectedCount(2)) == "reflected"
assert value == [1]


class RepeatMixin:
def __rmul__(self, value):
return "mixin"


class MixedRepeat(list, RepeatMixin):
pass


# A secondary base can supply the numeric slot even if list hides its method.
assert mul(MixedRepeat([1]), ReflectedCount(2)) == [1, 1]
assert imul(ReflectedCount(2), MixedRepeat([1])) == [1, 1]

# Removing an override preserves an already installed numeric dispatcher.
InheritedRepeat.__imul__ = lambda self, other: NotImplemented
del InheritedRepeat.__imul__
value = InheritedRepeat([1])
assert imul(value, ReflectedCount(2)) is value
assert value == [1, 1]


class ReflectedAdd(list):
def __radd__(self, value):
return "reflected"


InheritedRepeat.__iadd__ = lambda self, other: NotImplemented
del InheritedRepeat.__iadd__
value = InheritedRepeat([1])
assert iadd(value, ReflectedAdd([2])) is value
assert value == [1, 2]
Loading