Skip to content

Commit 31b89fc

Browse files
authored
Merge pull request RustPython#4347 from qingshi163/refactor-number
Refactor Number Protocol &PyNumber -> PyNumber
2 parents 1b6cb17 + f16a48d commit 31b89fc

File tree

6 files changed

+89
-109
lines changed

6 files changed

+89
-109
lines changed

stdlib/src/math.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -524,8 +524,8 @@ mod math {
524524
fn ceil(x: PyObjectRef, vm: &VirtualMachine) -> PyResult {
525525
let result_or_err = try_magic_method(identifier!(vm, __ceil__), vm, &x);
526526
if result_or_err.is_err() {
527-
if let Ok(Some(v)) = x.try_float_opt(vm) {
528-
let v = try_f64_to_bigint(v.to_f64().ceil(), vm)?;
527+
if let Some(v) = x.try_float_opt(vm) {
528+
let v = try_f64_to_bigint(v?.to_f64().ceil(), vm)?;
529529
return Ok(vm.ctx.new_int(v).into());
530530
}
531531
}
@@ -536,8 +536,8 @@ mod math {
536536
fn floor(x: PyObjectRef, vm: &VirtualMachine) -> PyResult {
537537
let result_or_err = try_magic_method(identifier!(vm, __floor__), vm, &x);
538538
if result_or_err.is_err() {
539-
if let Ok(Some(v)) = x.try_float_opt(vm) {
540-
let v = try_f64_to_bigint(v.to_f64().floor(), vm)?;
539+
if let Some(v) = x.try_float_opt(vm) {
540+
let v = try_f64_to_bigint(v?.to_f64().floor(), vm)?;
541541
return Ok(vm.ctx.new_int(v).into());
542542
}
543543
}

vm/src/builtins/complex.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,8 @@ impl PyObjectRef {
6969
if let Some(complex) = self.payload_if_subclass::<PyComplex>(vm) {
7070
return Ok(Some((complex.value, true)));
7171
}
72-
if let Some(float) = self.try_float_opt(vm)? {
73-
return Ok(Some((Complex64::new(float.to_f64(), 0.0), false)));
72+
if let Some(float) = self.try_float_opt(vm) {
73+
return Ok(Some((Complex64::new(float?.to_f64(), 0.0), false)));
7474
}
7575
Ok(None)
7676
}
@@ -489,7 +489,7 @@ impl AsNumber for PyComplex {
489489

490490
impl PyComplex {
491491
fn number_general_op<F, R>(
492-
number: &PyNumber,
492+
number: PyNumber,
493493
other: &PyObject,
494494
op: F,
495495
vm: &VirtualMachine,
@@ -506,7 +506,7 @@ impl PyComplex {
506506
}
507507

508508
fn number_complex_op<F>(
509-
number: &PyNumber,
509+
number: PyNumber,
510510
other: &PyObject,
511511
op: F,
512512
vm: &VirtualMachine,
@@ -517,7 +517,7 @@ impl PyComplex {
517517
Self::number_general_op(number, other, |a, b, _vm| op(a, b), vm)
518518
}
519519

520-
fn number_complex(number: &PyNumber, vm: &VirtualMachine) -> PyRef<PyComplex> {
520+
fn number_complex(number: PyNumber, vm: &VirtualMachine) -> PyRef<PyComplex> {
521521
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
522522
zelf.to_owned()
523523
} else {

vm/src/builtins/float.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -146,8 +146,8 @@ impl Constructor for PyFloat {
146146
return Ok(val);
147147
}
148148

149-
if let Some(f) = val.try_float_opt(vm)? {
150-
f.value
149+
if let Some(f) = val.try_float_opt(vm) {
150+
f?.value
151151
} else {
152152
float_from_string(val, vm)?
153153
}
@@ -608,7 +608,7 @@ impl AsNumber for PyFloat {
608608

609609
impl PyFloat {
610610
fn number_general_op<F, R>(
611-
number: &PyNumber,
611+
number: PyNumber,
612612
other: &PyObject,
613613
op: F,
614614
vm: &VirtualMachine,
@@ -625,7 +625,7 @@ impl PyFloat {
625625
}
626626

627627
fn number_float_op<F>(
628-
number: &PyNumber,
628+
number: PyNumber,
629629
other: &PyObject,
630630
op: F,
631631
vm: &VirtualMachine,
@@ -636,7 +636,7 @@ impl PyFloat {
636636
Self::number_general_op(number, other, |a, b, _vm| op(a, b), vm)
637637
}
638638

639-
fn number_float(number: &PyNumber, vm: &VirtualMachine) -> PyRef<PyFloat> {
639+
fn number_float(number: PyNumber, vm: &VirtualMachine) -> PyRef<PyFloat> {
640640
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
641641
zelf.to_owned()
642642
} else {

vm/src/builtins/int.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -802,7 +802,7 @@ impl AsNumber for PyInt {
802802

803803
impl PyInt {
804804
fn number_general_op<F>(
805-
number: &PyNumber,
805+
number: PyNumber,
806806
other: &PyObject,
807807
op: F,
808808
vm: &VirtualMachine,
@@ -817,14 +817,14 @@ impl PyInt {
817817
}
818818
}
819819

820-
fn number_int_op<F>(number: &PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
820+
fn number_int_op<F>(number: PyNumber, other: &PyObject, op: F, vm: &VirtualMachine) -> PyResult
821821
where
822822
F: FnOnce(&BigInt, &BigInt) -> BigInt,
823823
{
824824
Self::number_general_op(number, other, |a, b, _vm| op(a, b).to_pyresult(vm), vm)
825825
}
826826

827-
fn number_int(number: &PyNumber, vm: &VirtualMachine) -> PyIntRef {
827+
fn number_int(number: PyNumber, vm: &VirtualMachine) -> PyIntRef {
828828
if let Some(zelf) = number.obj.downcast_ref_if_exact::<Self>(vm) {
829829
zelf.to_owned()
830830
} else {

vm/src/protocol/number.rs

Lines changed: 67 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,9 @@ use crate::{
99
};
1010
use crossbeam_utils::atomic::AtomicCell;
1111

12-
type UnaryFunc<R = PyObjectRef> = AtomicCell<Option<fn(&PyNumber, &VirtualMachine) -> PyResult<R>>>;
12+
type UnaryFunc<R = PyObjectRef> = AtomicCell<Option<fn(PyNumber, &VirtualMachine) -> PyResult<R>>>;
1313
type BinaryFunc<R = PyObjectRef> =
14-
AtomicCell<Option<fn(&PyNumber, &PyObject, &VirtualMachine) -> PyResult<R>>>;
14+
AtomicCell<Option<fn(PyNumber, &PyObject, &VirtualMachine) -> PyResult<R>>>;
1515

1616
impl PyObject {
1717
#[inline]
@@ -20,16 +20,13 @@ impl PyObject {
2020
}
2121

2222
pub fn try_index_opt(&self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
23-
#[allow(clippy::question_mark)]
24-
Some(if let Some(i) = self.downcast_ref_if_exact::<PyInt>(vm) {
25-
Ok(i.to_owned())
23+
if let Some(i) = self.downcast_ref_if_exact::<PyInt>(vm) {
24+
Some(Ok(i.to_owned()))
2625
} else if let Some(i) = self.payload::<PyInt>() {
27-
Ok(vm.ctx.new_bigint(i.as_bigint()))
28-
} else if let Some(i) = self.to_number().index(vm).transpose() {
29-
i
26+
Some(Ok(vm.ctx.new_bigint(i.as_bigint())))
3027
} else {
31-
return None;
32-
})
28+
self.to_number().index(vm)
29+
}
3330
}
3431

3532
#[inline]
@@ -57,71 +54,57 @@ impl PyObject {
5754

5855
if let Some(i) = self.downcast_ref_if_exact::<PyInt>(vm) {
5956
Ok(i.to_owned())
57+
} else if let Some(i) = self.to_number().int(vm).or_else(|| self.try_index_opt(vm)) {
58+
i
59+
} else if let Ok(Ok(f)) = vm.get_special_method(self.to_owned(), identifier!(vm, __trunc__))
60+
{
61+
// TODO: Deprecate in 3.11
62+
// warnings::warn(
63+
// vm.ctx.exceptions.deprecation_warning.clone(),
64+
// "The delegation of int() to __trunc__ is deprecated.".to_owned(),
65+
// 1,
66+
// vm,
67+
// )?;
68+
let ret = f.invoke((), vm)?;
69+
ret.try_index(vm).map_err(|_| {
70+
vm.new_type_error(format!(
71+
"__trunc__ returned non-Integral (type {})",
72+
ret.class()
73+
))
74+
})
75+
} else if let Some(s) = self.payload::<PyStr>() {
76+
try_convert(self, s.as_str().as_bytes(), vm)
77+
} else if let Some(bytes) = self.payload::<PyBytes>() {
78+
try_convert(self, bytes, vm)
79+
} else if let Some(bytearray) = self.payload::<PyByteArray>() {
80+
try_convert(self, &bytearray.borrow_buf(), vm)
81+
} else if let Ok(buffer) = ArgBytesLike::try_from_borrowed_object(vm, self) {
82+
// TODO: replace to PyBuffer
83+
try_convert(self, &buffer.borrow_buf(), vm)
6084
} else {
61-
let number = self.to_number();
62-
if let Some(i) = number.int(vm)? {
63-
Ok(i)
64-
} else if let Some(i) = self.try_index_opt(vm) {
65-
i
66-
} else if let Ok(Ok(f)) =
67-
vm.get_special_method(self.to_owned(), identifier!(vm, __trunc__))
68-
{
69-
// TODO: Deprecate in 3.11
70-
// warnings::warn(
71-
// vm.ctx.exceptions.deprecation_warning.clone(),
72-
// "The delegation of int() to __trunc__ is deprecated.".to_owned(),
73-
// 1,
74-
// vm,
75-
// )?;
76-
let ret = f.invoke((), vm)?;
77-
ret.try_index(vm).map_err(|_| {
78-
vm.new_type_error(format!(
79-
"__trunc__ returned non-Integral (type {})",
80-
ret.class()
81-
))
82-
})
83-
} else if let Some(s) = self.payload::<PyStr>() {
84-
try_convert(self, s.as_str().as_bytes(), vm)
85-
} else if let Some(bytes) = self.payload::<PyBytes>() {
86-
try_convert(self, bytes, vm)
87-
} else if let Some(bytearray) = self.payload::<PyByteArray>() {
88-
try_convert(self, &bytearray.borrow_buf(), vm)
89-
} else if let Ok(buffer) = ArgBytesLike::try_from_borrowed_object(vm, self) {
90-
// TODO: replace to PyBuffer
91-
try_convert(self, &buffer.borrow_buf(), vm)
92-
} else {
93-
Err(vm.new_type_error(format!(
94-
"int() argument must be a string, a bytes-like object or a real number, not '{}'",
95-
self.class()
96-
)))
97-
}
85+
Err(vm.new_type_error(format!(
86+
"int() argument must be a string, a bytes-like object or a real number, not '{}'",
87+
self.class()
88+
)))
9889
}
9990
}
10091

101-
pub fn try_float_opt(&self, vm: &VirtualMachine) -> PyResult<Option<PyRef<PyFloat>>> {
102-
let value = if let Some(float) = self.downcast_ref_if_exact::<PyFloat>(vm) {
103-
Some(float.to_owned())
92+
pub fn try_float_opt(&self, vm: &VirtualMachine) -> Option<PyResult<PyRef<PyFloat>>> {
93+
if let Some(float) = self.downcast_ref_if_exact::<PyFloat>(vm) {
94+
Some(Ok(float.to_owned()))
95+
} else if let Some(f) = self.to_number().float(vm) {
96+
Some(f)
10497
} else {
105-
let number = self.to_number();
106-
#[allow(clippy::manual_map)]
107-
if let Some(f) = number.float(vm)? {
108-
Some(f)
109-
} else if let Some(i) = self.try_index_opt(vm) {
110-
let value = int::try_to_float(i?.as_bigint(), vm)?;
111-
Some(vm.ctx.new_float(value))
112-
} else if let Some(value) = self.downcast_ref::<PyFloat>() {
113-
Some(vm.ctx.new_float(value.to_f64()))
114-
} else {
115-
None
116-
}
117-
};
118-
Ok(value)
98+
self.try_index_opt(vm)
99+
.map(|i| Ok(vm.ctx.new_float(int::try_to_float(i?.as_bigint(), vm)?)))
100+
}
119101
}
120102

121103
#[inline]
122104
pub fn try_float(&self, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
123-
self.try_float_opt(vm)?
124-
.ok_or_else(|| vm.new_type_error(format!("must be real number, not {}", self.class())))
105+
self.try_float_opt(vm).ok_or_else(|| {
106+
vm.new_type_error(format!("must be real number, not {}", self.class()))
107+
})?
125108
}
126109
}
127110

@@ -259,10 +242,10 @@ impl PyNumber<'_> {
259242
}
260243

261244
#[inline]
262-
pub fn int(&self, vm: &VirtualMachine) -> PyResult<Option<PyIntRef>> {
263-
Ok(if let Some(f) = self.methods().int.load() {
245+
pub fn int(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
246+
self.methods().int.load().map(|f| {
264247
let ret = f(self, vm)?;
265-
Some(if !ret.class().is(PyInt::class(vm)) {
248+
let value = if !ret.class().is(PyInt::class(vm)) {
266249
warnings::warn(
267250
vm.ctx.exceptions.deprecation_warning,
268251
format!(
@@ -277,17 +260,16 @@ impl PyNumber<'_> {
277260
vm.ctx.new_bigint(ret.as_bigint())
278261
} else {
279262
ret
280-
})
281-
} else {
282-
None
263+
};
264+
Ok(value)
283265
})
284266
}
285267

286268
#[inline]
287-
pub fn index(&self, vm: &VirtualMachine) -> PyResult<Option<PyIntRef>> {
288-
if let Some(f) = self.methods().index.load() {
269+
pub fn index(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
270+
self.methods().index.load().map(|f| {
289271
let ret = f(self, vm)?;
290-
if !ret.class().is(PyInt::class(vm)) {
272+
let value = if !ret.class().is(PyInt::class(vm)) {
291273
warnings::warn(
292274
vm.ctx.exceptions.deprecation_warning,
293275
format!(
@@ -299,20 +281,19 @@ impl PyNumber<'_> {
299281
1,
300282
vm,
301283
)?;
302-
Ok(Some(vm.ctx.new_bigint(ret.as_bigint())))
284+
vm.ctx.new_bigint(ret.as_bigint())
303285
} else {
304-
Ok(Some(ret))
305-
}
306-
} else {
307-
Ok(None)
308-
}
286+
ret
287+
};
288+
Ok(value)
289+
})
309290
}
310291

311292
#[inline]
312-
pub fn float(&self, vm: &VirtualMachine) -> PyResult<Option<PyRef<PyFloat>>> {
313-
Ok(if let Some(f) = self.methods().float.load() {
293+
pub fn float(self, vm: &VirtualMachine) -> Option<PyResult<PyRef<PyFloat>>> {
294+
self.methods().float.load().map(|f| {
314295
let ret = f(self, vm)?;
315-
Some(if !ret.class().is(PyFloat::class(vm)) {
296+
let value = if !ret.class().is(PyFloat::class(vm)) {
316297
warnings::warn(
317298
vm.ctx.exceptions.deprecation_warning,
318299
format!(
@@ -327,9 +308,8 @@ impl PyNumber<'_> {
327308
vm.ctx.new_float(ret.to_f64())
328309
} else {
329310
ret
330-
})
331-
} else {
332-
None
311+
};
312+
Ok(value)
333313
})
334314
}
335315
}

vm/src/types/slot.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -190,21 +190,21 @@ pub(crate) fn len_wrapper(obj: &PyObject, vm: &VirtualMachine) -> PyResult<usize
190190
Ok(len as usize)
191191
}
192192

193-
fn int_wrapper(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
193+
fn int_wrapper(num: PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
194194
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __int__), ())?;
195195
ret.downcast::<PyInt>().map_err(|obj| {
196196
vm.new_type_error(format!("__int__ returned non-int (type {})", obj.class()))
197197
})
198198
}
199199

200-
fn index_wrapper(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
200+
fn index_wrapper(num: PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyInt>> {
201201
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __index__), ())?;
202202
ret.downcast::<PyInt>().map_err(|obj| {
203203
vm.new_type_error(format!("__index__ returned non-int (type {})", obj.class()))
204204
})
205205
}
206206

207-
fn float_wrapper(num: &PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
207+
fn float_wrapper(num: PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
208208
let ret = vm.call_special_method(num.obj.to_owned(), identifier!(vm, __float__), ())?;
209209
ret.downcast::<PyFloat>().map_err(|obj| {
210210
vm.new_type_error(format!(
@@ -968,8 +968,8 @@ pub trait AsNumber: PyPayload {
968968
fn as_number() -> &'static PyNumberMethods;
969969

970970
#[inline]
971-
fn number_downcast<'a>(number: &'a PyNumber) -> &'a Py<Self> {
972-
unsafe { number.obj.downcast_unchecked_ref() }
971+
fn number_downcast(num: PyNumber) -> &Py<Self> {
972+
unsafe { num.obj.downcast_unchecked_ref() }
973973
}
974974
}
975975

0 commit comments

Comments
 (0)