Skip to content

Commit 21c7e15

Browse files
YashSuthar983YashSuthar983
committed
Resolve number slots via MRO in PyNumber and operator, ensure inherited and dynamically added methods are found.
Use class().mro_find_map() to mimic the same behaviour as CPython. Signed-off-by: Yash Suthar <yashsuthar983@gmail.com>
1 parent 9825d8a commit 21c7e15

File tree

2 files changed

+120
-94
lines changed

2 files changed

+120
-94
lines changed

‎vm/src/protocol/number.rs‎

Lines changed: 105 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -443,117 +443,135 @@ impl<'a> PyNumber<'a> {
443443

444444
// PyNumber_Check
445445
pub fn check(obj: &PyObject) -> bool {
446-
let methods = &obj.class().slots.as_number;
447-
methods.int.load().is_some()
448-
|| methods.index.load().is_some()
449-
|| methods.float.load().is_some()
450-
|| obj.downcastable::<PyComplex>()
446+
let cls = &obj.class();
447+
let has_number = cls
448+
.mro_find_map(|x| {
449+
let methods = &x.slots.as_number;
450+
if methods.int.load().is_some()
451+
|| methods.index.load().is_some()
452+
|| methods.float.load().is_some()
453+
{
454+
Some(())
455+
} else {
456+
None
457+
}
458+
})
459+
.is_some();
460+
has_number || obj.downcastable::<PyComplex>()
451461
}
452462
}
453463

454464
impl PyNumber<'_> {
455465
// PyIndex_Check
456466
pub fn is_index(self) -> bool {
457-
self.class().slots.as_number.index.load().is_some()
467+
self.class()
468+
.mro_find_map(|x| x.slots.as_number.index.load())
469+
.is_some()
458470
}
459471

460472
#[inline]
461473
pub fn int(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
462-
self.class().slots.as_number.int.load().map(|f| {
463-
let ret = f(self, vm)?;
464-
465-
if let Some(ret) = ret.downcast_ref_if_exact::<PyInt>(vm) {
466-
return Ok(ret.to_owned());
467-
}
468-
469-
let ret_class = ret.class().to_owned();
470-
if let Some(ret) = ret.downcast_ref::<PyInt>() {
471-
warnings::warn(
472-
vm.ctx.exceptions.deprecation_warning,
473-
format!(
474-
"__int__ returned non-int (type {ret_class}). \
474+
self.class()
475+
.mro_find_map(|x| x.slots.as_number.int.load())
476+
.map(|f| {
477+
let ret = f(self, vm)?;
478+
479+
if let Some(ret) = ret.downcast_ref_if_exact::<PyInt>(vm) {
480+
return Ok(ret.to_owned());
481+
}
482+
483+
let ret_class = ret.class().to_owned();
484+
if let Some(ret) = ret.downcast_ref::<PyInt>() {
485+
warnings::warn(
486+
vm.ctx.exceptions.deprecation_warning,
487+
format!(
488+
"__int__ returned non-int (type {ret_class}). \
475489
The ability to return an instance of a strict subclass of int \
476490
is deprecated, and may be removed in a future version of Python."
477-
),
478-
1,
479-
vm,
480-
)?;
481-
482-
Ok(ret.to_owned())
483-
} else {
484-
Err(vm.new_type_error(format!(
485-
"{}.__int__ returned non-int(type {})",
486-
self.class(),
487-
ret_class
488-
)))
489-
}
490-
})
491+
),
492+
1,
493+
vm,
494+
)?;
495+
496+
Ok(ret.to_owned())
497+
} else {
498+
Err(vm.new_type_error(format!(
499+
"{}.__int__ returned non-int(type {})",
500+
self.class(),
501+
ret_class
502+
)))
503+
}
504+
})
491505
}
492506

493507
#[inline]
494508
pub fn index(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
495-
self.class().slots.as_number.index.load().map(|f| {
496-
let ret = f(self, vm)?;
497-
498-
if let Some(ret) = ret.downcast_ref_if_exact::<PyInt>(vm) {
499-
return Ok(ret.to_owned());
500-
}
501-
502-
let ret_class = ret.class().to_owned();
503-
if let Some(ret) = ret.downcast_ref::<PyInt>() {
504-
warnings::warn(
505-
vm.ctx.exceptions.deprecation_warning,
506-
format!(
507-
"__index__ returned non-int (type {ret_class}). \
509+
self.class()
510+
.mro_find_map(|x| x.slots.as_number.index.load())
511+
.map(|f| {
512+
let ret = f(self, vm)?;
513+
514+
if let Some(ret) = ret.downcast_ref_if_exact::<PyInt>(vm) {
515+
return Ok(ret.to_owned());
516+
}
517+
518+
let ret_class = ret.class().to_owned();
519+
if let Some(ret) = ret.downcast_ref::<PyInt>() {
520+
warnings::warn(
521+
vm.ctx.exceptions.deprecation_warning,
522+
format!(
523+
"__index__ returned non-int (type {ret_class}). \
508524
The ability to return an instance of a strict subclass of int \
509525
is deprecated, and may be removed in a future version of Python."
510-
),
511-
1,
512-
vm,
513-
)?;
514-
515-
Ok(ret.to_owned())
516-
} else {
517-
Err(vm.new_type_error(format!(
518-
"{}.__index__ returned non-int(type {})",
519-
self.class(),
520-
ret_class
521-
)))
522-
}
523-
})
526+
),
527+
1,
528+
vm,
529+
)?;
530+
531+
Ok(ret.to_owned())
532+
} else {
533+
Err(vm.new_type_error(format!(
534+
"{}.__index__ returned non-int(type {})",
535+
self.class(),
536+
ret_class
537+
)))
538+
}
539+
})
524540
}
525541

526542
#[inline]
527543
pub fn float(self, vm: &VirtualMachine) -> Option<PyResult<PyRef<PyFloat>>> {
528-
self.class().slots.as_number.float.load().map(|f| {
529-
let ret = f(self, vm)?;
530-
531-
if let Some(ret) = ret.downcast_ref_if_exact::<PyFloat>(vm) {
532-
return Ok(ret.to_owned());
533-
}
534-
535-
let ret_class = ret.class().to_owned();
536-
if let Some(ret) = ret.downcast_ref::<PyFloat>() {
537-
warnings::warn(
538-
vm.ctx.exceptions.deprecation_warning,
539-
format!(
540-
"__float__ returned non-float (type {ret_class}). \
544+
self.class()
545+
.mro_find_map(|x| x.slots.as_number.float.load())
546+
.map(|f| {
547+
let ret = f(self, vm)?;
548+
549+
if let Some(ret) = ret.downcast_ref_if_exact::<PyFloat>(vm) {
550+
return Ok(ret.to_owned());
551+
}
552+
553+
let ret_class = ret.class().to_owned();
554+
if let Some(ret) = ret.downcast_ref::<PyFloat>() {
555+
warnings::warn(
556+
vm.ctx.exceptions.deprecation_warning,
557+
format!(
558+
"__float__ returned non-float (type {ret_class}). \
541559
The ability to return an instance of a strict subclass of float \
542560
is deprecated, and may be removed in a future version of Python."
543-
),
544-
1,
545-
vm,
546-
)?;
547-
548-
Ok(ret.to_owned())
549-
} else {
550-
Err(vm.new_type_error(format!(
551-
"{}.__float__ returned non-float(type {})",
552-
self.class(),
553-
ret_class
554-
)))
555-
}
556-
})
561+
),
562+
1,
563+
vm,
564+
)?;
565+
566+
Ok(ret.to_owned())
567+
} else {
568+
Err(vm.new_type_error(format!(
569+
"{}.__float__ returned non-float(type {})",
570+
self.class(),
571+
ret_class
572+
)))
573+
}
574+
})
557575
}
558576
}
559577

‎vm/src/vm/vm_ops.rs‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -160,11 +160,12 @@ impl VirtualMachine {
160160
let class_a = a.class();
161161
let class_b = b.class();
162162

163-
let slot_a = class_a.slots.as_number.left_binary_op(op_slot);
163+
// Look up number slots across MRO for inheritance
164+
let slot_a = class_a.mro_find_map(|x| x.slots.as_number.left_binary_op(op_slot));
164165
let mut slot_b = None;
165166

166167
if !class_a.is(class_b) {
167-
let slot_bb = class_b.slots.as_number.right_binary_op(op_slot);
168+
let slot_bb = class_b.mro_find_map(|x| x.slots.as_number.right_binary_op(op_slot));
168169
if slot_bb.map(|x| x as usize) != slot_a.map(|x| x as usize) {
169170
slot_b = slot_bb;
170171
}
@@ -230,7 +231,10 @@ impl VirtualMachine {
230231
iop_slot: PyNumberBinaryOp,
231232
op_slot: PyNumberBinaryOp,
232233
) -> PyResult {
233-
if let Some(slot) = a.class().slots.as_number.left_binary_op(iop_slot) {
234+
if let Some(slot) = a
235+
.class()
236+
.mro_find_map(|x| x.slots.as_number.left_binary_op(iop_slot))
237+
{
234238
let x = slot(a, b, self)?;
235239
if !x.is(&self.ctx.not_implemented) {
236240
return Ok(x);
@@ -266,11 +270,12 @@ impl VirtualMachine {
266270
let class_b = b.class();
267271
let class_c = c.class();
268272

269-
let slot_a = class_a.slots.as_number.left_ternary_op(op_slot);
273+
// Look up number slots across MRO for inheritance
274+
let slot_a = class_a.mro_find_map(|x| x.slots.as_number.left_ternary_op(op_slot));
270275
let mut slot_b = None;
271276

272277
if !class_a.is(class_b) {
273-
let slot_bb = class_b.slots.as_number.right_ternary_op(op_slot);
278+
let slot_bb = class_b.mro_find_map(|x| x.slots.as_number.right_ternary_op(op_slot));
274279
if slot_bb.map(|x| x as usize) != slot_a.map(|x| x as usize) {
275280
slot_b = slot_bb;
276281
}
@@ -299,7 +304,7 @@ impl VirtualMachine {
299304
}
300305
}
301306

302-
if let Some(slot_c) = class_c.slots.as_number.left_ternary_op(op_slot)
307+
if let Some(slot_c) = class_c.mro_find_map(|x| x.slots.as_number.left_ternary_op(op_slot))
303308
&& slot_a.is_some_and(|slot_a| !std::ptr::fn_addr_eq(slot_a, slot_c))
304309
&& slot_b.is_some_and(|slot_b| !std::ptr::fn_addr_eq(slot_b, slot_c))
305310
{
@@ -338,7 +343,10 @@ impl VirtualMachine {
338343
op_slot: PyNumberTernaryOp,
339344
op_str: &str,
340345
) -> PyResult {
341-
if let Some(slot) = a.class().slots.as_number.left_ternary_op(iop_slot) {
346+
if let Some(slot) = a
347+
.class()
348+
.mro_find_map(|x| x.slots.as_number.left_ternary_op(iop_slot))
349+
{
342350
let x = slot(a, b, c, self)?;
343351
if !x.is(&self.ctx.not_implemented) {
344352
return Ok(x);

0 commit comments

Comments
 (0)