Skip to content

Commit beeb449

Browse files
qingshi163youknowone
qingshi163
authored and
youknowone
committed
PyNumber protocol
1 parent 8e20e23 commit beeb449

File tree

3 files changed

+164
-1
lines changed

3 files changed

+164
-1
lines changed

‎vm/src/builtins/int.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -807,7 +807,7 @@ fn try_int_radix(obj: &PyObject, base: u32, vm: &VirtualMachine) -> PyResult<Big
807807
}
808808
}
809809

810-
fn bytes_to_int(lit: &[u8], mut base: u32) -> Option<BigInt> {
810+
pub(crate) fn bytes_to_int(lit: &[u8], mut base: u32) -> Option<BigInt> {
811811
// split sign
812812
let mut lit = lit.trim();
813813
let sign = match lit.first()? {

‎vm/src/protocol/number.rs‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
use std::borrow::Cow;
2+
3+
use crate::{
4+
builtins::{int, PyByteArray, PyBytes, PyComplex, PyFloat, PyInt, PyIntRef, PyStr},
5+
common::{lock::OnceCell, static_cell},
6+
function::ArgBytesLike,
7+
IdProtocol, PyObject, PyRef, PyResult, PyValue, TryFromBorrowedObject, TypeProtocol,
8+
VirtualMachine,
9+
};
10+
11+
#[allow(clippy::type_complexity)]
12+
#[derive(Default, Clone)]
13+
pub struct PyNumberMethods {
14+
/* Number implementations must check *both*
15+
arguments for proper type and implement the necessary conversions
16+
in the slot functions themselves. */
17+
pub add: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
18+
pub subtract: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
19+
pub multiply: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
20+
pub remainder: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
21+
pub divmod: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
22+
pub power: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
23+
pub negative: Option<fn(&PyNumber, vm: &VirtualMachine) -> PyResult>,
24+
pub positive: Option<fn(&PyNumber, vm: &VirtualMachine) -> PyResult>,
25+
pub absolute: Option<fn(&PyNumber, vm: &VirtualMachine) -> PyResult>,
26+
pub boolean: Option<fn(&PyNumber, vm: &VirtualMachine) -> PyResult<bool>>,
27+
pub invert: Option<fn(&PyNumber, vm: &VirtualMachine) -> PyResult>,
28+
pub lshift: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
29+
pub rshift: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
30+
pub and: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
31+
pub xor: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
32+
pub or: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
33+
pub int: Option<fn(&PyNumber, vm: &VirtualMachine) -> PyResult<PyIntRef>>,
34+
pub float: Option<fn(&PyNumber, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>>>,
35+
36+
pub inplace_add: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
37+
pub inplace_substract: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
38+
pub inplace_multiply: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
39+
pub inplace_remainder: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
40+
pub inplace_divmod: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
41+
pub inplace_power: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
42+
pub inplace_lshift: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
43+
pub inplace_rshift: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
44+
pub inplace_and: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
45+
pub inplace_xor: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
46+
pub inplace_or: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
47+
48+
pub floor_divide: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
49+
pub true_divide: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
50+
pub inplace_floor_divide: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
51+
pub inplace_true_devide: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
52+
53+
pub index: Option<fn(&PyNumber, vm: &VirtualMachine) -> PyResult<PyIntRef>>,
54+
55+
pub matrix_multiply: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
56+
pub inplace_matrix_multiply: Option<fn(&PyNumber, &PyObject, vm: &VirtualMachine) -> PyResult>,
57+
}
58+
59+
impl PyNumberMethods {
60+
fn not_implemented() -> &'static Self {
61+
static_cell! {
62+
static NOT_IMPLEMENTED: PyNumberMethods;
63+
}
64+
NOT_IMPLEMENTED.get_or_init(Self::default)
65+
}
66+
}
67+
68+
pub struct PyNumber<'a> {
69+
pub obj: &'a PyObject,
70+
// some fast path do not need methods, so we do lazy initialize
71+
methods: OnceCell<Cow<'static, PyNumberMethods>>,
72+
}
73+
74+
impl<'a> From<&'a PyObject> for PyNumber<'a> {
75+
fn from(obj: &'a PyObject) -> Self {
76+
Self {
77+
obj,
78+
methods: OnceCell::new(),
79+
}
80+
}
81+
}
82+
83+
impl<'a> PyNumber<'a> {
84+
pub fn methods(&'a self, vm: &VirtualMachine) -> &'a Cow<'static, PyNumberMethods> {
85+
self.methods.get_or_init(|| {
86+
self.obj
87+
.class()
88+
.mro_find_map(|x| x.slots.as_number.load())
89+
.map(|f| f(self.obj, vm))
90+
.unwrap_or_else(|| Cow::Borrowed(PyNumberMethods::not_implemented()))
91+
})
92+
}
93+
}
94+
95+
impl PyNumber<'_> {
96+
// PyNumber_Check
97+
pub fn is_numeric(&self, vm: &VirtualMachine) -> bool {
98+
let methods = self.methods(vm);
99+
methods.int.is_some()
100+
|| methods.index.is_some()
101+
|| methods.float.is_some()
102+
|| self.obj.payload_is::<PyComplex>()
103+
}
104+
105+
// PyIndex_Check
106+
pub fn is_index(&self, vm: &VirtualMachine) -> bool {
107+
self.methods(vm).index.is_some()
108+
}
109+
110+
pub fn to_int(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
111+
fn try_convert(obj: &PyObject, lit: &[u8], vm: &VirtualMachine) -> PyResult<PyIntRef> {
112+
let base = 10;
113+
match int::bytes_to_int(lit, base) {
114+
Some(i) => Ok(PyInt::from(i).into_ref(vm)),
115+
None => Err(vm.new_value_error(format!(
116+
"invalid literal for int() with base {}: {}",
117+
base,
118+
obj.repr(vm)?,
119+
))),
120+
}
121+
}
122+
123+
if self.obj.class().is(PyInt::class(vm)) {
124+
Ok(unsafe { self.obj.downcast_unchecked_ref::<PyInt>() }.to_owned())
125+
} else if let Some(f) = self.methods(vm).int {
126+
f(self, vm)
127+
} else if let Some(f) = self.methods(vm).index {
128+
f(self, vm)
129+
} else if let Ok(Ok(f)) = vm.get_special_method(self.obj.to_owned(), "__trunc__") {
130+
let r = f.invoke((), vm)?;
131+
PyNumber::from(r.as_ref()).to_index(vm)
132+
} else if let Some(s) = self.obj.payload::<PyStr>() {
133+
try_convert(self.obj, s.as_str().as_bytes(), vm)
134+
} else if let Some(bytes) = self.obj.payload::<PyBytes>() {
135+
try_convert(self.obj, bytes, vm)
136+
} else if let Some(bytearray) = self.obj.payload::<PyByteArray>() {
137+
try_convert(self.obj, &bytearray.borrow_buf(), vm)
138+
} else if let Ok(buffer) = ArgBytesLike::try_from_borrowed_object(vm, self.obj) {
139+
// TODO: replace to PyBuffer
140+
try_convert(self.obj, &buffer.borrow_buf(), vm)
141+
} else {
142+
Err(vm.new_type_error(format!(
143+
"int() argument must be a string, a bytes-like object or a real number, not '{}'",
144+
self.obj.class()
145+
)))
146+
}
147+
}
148+
149+
pub fn to_index(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
150+
if self.obj.class().is(PyInt::class(vm)) {
151+
Ok(unsafe { self.obj.downcast_unchecked_ref::<PyInt>() }.to_owned())
152+
} else if let Some(f) = self.methods(vm).index {
153+
f(self, vm)
154+
} else {
155+
Err(vm.new_type_error(format!(
156+
"'{}' object cannot be interpreted as an integer",
157+
self.obj.class()
158+
)))
159+
}
160+
}
161+
}

‎vm/src/types/slot.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use crate::common::{hash::PyHash, lock::PyRwLock};
2+
use crate::protocol::PyNumberMethods;
23
use crate::{
34
builtins::{PyInt, PyStrInterned, PyStrRef, PyType, PyTypeRef},
45
bytecode::ComparisonOperator,
@@ -138,6 +139,7 @@ impl Default for PyTypeFlags {
138139

139140
pub(crate) type GenericMethod = fn(&PyObject, FuncArgs, &VirtualMachine) -> PyResult;
140141
pub(crate) type AsMappingFunc = fn(&PyObject, &VirtualMachine) -> &'static PyMappingMethods;
142+
pub(crate) type AsNumberFunc = fn(&PyObject, &VirtualMachine) -> Cow<'static, PyNumberMethods>;
141143
pub(crate) type HashFunc = fn(&PyObject, &VirtualMachine) -> PyResult<PyHash>;
142144
// CallFunc = GenericMethod
143145
pub(crate) type GetattroFunc = fn(&PyObject, PyStrRef, &VirtualMachine) -> PyResult;

0 commit comments

Comments
 (0)