rustpython_vm/protocol/
number.rs

1use std::ops::Deref;
2
3use crossbeam_utils::atomic::AtomicCell;
4
5use crate::{
6    builtins::{int, PyByteArray, PyBytes, PyComplex, PyFloat, PyInt, PyIntRef, PyStr},
7    common::int::bytes_to_int,
8    function::ArgBytesLike,
9    object::{Traverse, TraverseFn},
10    stdlib::warnings,
11    AsObject, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromBorrowedObject,
12    VirtualMachine,
13};
14
15pub type PyNumberUnaryFunc<R = PyObjectRef> = fn(PyNumber, &VirtualMachine) -> PyResult<R>;
16pub type PyNumberBinaryFunc = fn(&PyObject, &PyObject, &VirtualMachine) -> PyResult;
17pub type PyNumberTernaryFunc = fn(&PyObject, &PyObject, &PyObject, &VirtualMachine) -> PyResult;
18
19impl PyObject {
20    #[inline]
21    pub fn to_number(&self) -> PyNumber {
22        PyNumber(self)
23    }
24
25    pub fn try_index_opt(&self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
26        if let Some(i) = self.downcast_ref_if_exact::<PyInt>(vm) {
27            Some(Ok(i.to_owned()))
28        } else if let Some(i) = self.payload::<PyInt>() {
29            Some(Ok(vm.ctx.new_bigint(i.as_bigint())))
30        } else {
31            self.to_number().index(vm)
32        }
33    }
34
35    #[inline]
36    pub fn try_index(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
37        self.try_index_opt(vm).transpose()?.ok_or_else(|| {
38            vm.new_type_error(format!(
39                "'{}' object cannot be interpreted as an integer",
40                self.class()
41            ))
42        })
43    }
44
45    pub fn try_int(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
46        fn try_convert(obj: &PyObject, lit: &[u8], vm: &VirtualMachine) -> PyResult<PyIntRef> {
47            let base = 10;
48            let i = bytes_to_int(lit, base).ok_or_else(|| {
49                let repr = match obj.repr(vm) {
50                    Ok(repr) => repr,
51                    Err(err) => return err,
52                };
53                vm.new_value_error(format!(
54                    "invalid literal for int() with base {}: {}",
55                    base, repr,
56                ))
57            })?;
58            Ok(PyInt::from(i).into_ref(&vm.ctx))
59        }
60
61        if let Some(i) = self.downcast_ref_if_exact::<PyInt>(vm) {
62            Ok(i.to_owned())
63        } else if let Some(i) = self.to_number().int(vm).or_else(|| self.try_index_opt(vm)) {
64            i
65        } else if let Ok(Some(f)) = vm.get_special_method(self, identifier!(vm, __trunc__)) {
66            // TODO: Deprecate in 3.11
67            // warnings::warn(
68            //     vm.ctx.exceptions.deprecation_warning.clone(),
69            //     "The delegation of int() to __trunc__ is deprecated.".to_owned(),
70            //     1,
71            //     vm,
72            // )?;
73            let ret = f.invoke((), vm)?;
74            ret.try_index(vm).map_err(|_| {
75                vm.new_type_error(format!(
76                    "__trunc__ returned non-Integral (type {})",
77                    ret.class()
78                ))
79            })
80        } else if let Some(s) = self.payload::<PyStr>() {
81            try_convert(self, s.as_str().as_bytes(), vm)
82        } else if let Some(bytes) = self.payload::<PyBytes>() {
83            try_convert(self, bytes, vm)
84        } else if let Some(bytearray) = self.payload::<PyByteArray>() {
85            try_convert(self, &bytearray.borrow_buf(), vm)
86        } else if let Ok(buffer) = ArgBytesLike::try_from_borrowed_object(vm, self) {
87            // TODO: replace to PyBuffer
88            try_convert(self, &buffer.borrow_buf(), vm)
89        } else {
90            Err(vm.new_type_error(format!(
91                "int() argument must be a string, a bytes-like object or a real number, not '{}'",
92                self.class()
93            )))
94        }
95    }
96
97    pub fn try_float_opt(&self, vm: &VirtualMachine) -> Option<PyResult<PyRef<PyFloat>>> {
98        if let Some(float) = self.downcast_ref_if_exact::<PyFloat>(vm) {
99            Some(Ok(float.to_owned()))
100        } else if let Some(f) = self.to_number().float(vm) {
101            Some(f)
102        } else {
103            self.try_index_opt(vm)
104                .map(|i| Ok(vm.ctx.new_float(int::try_to_float(i?.as_bigint(), vm)?)))
105        }
106    }
107
108    #[inline]
109    pub fn try_float(&self, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
110        self.try_float_opt(vm).ok_or_else(|| {
111            vm.new_type_error(format!("must be real number, not {}", self.class()))
112        })?
113    }
114}
115
116#[derive(Default)]
117pub struct PyNumberMethods {
118    /* Number implementations must check *both*
119    arguments for proper type and implement the necessary conversions
120    in the slot functions themselves. */
121    pub add: Option<PyNumberBinaryFunc>,
122    pub subtract: Option<PyNumberBinaryFunc>,
123    pub multiply: Option<PyNumberBinaryFunc>,
124    pub remainder: Option<PyNumberBinaryFunc>,
125    pub divmod: Option<PyNumberBinaryFunc>,
126    pub power: Option<PyNumberTernaryFunc>,
127    pub negative: Option<PyNumberUnaryFunc>,
128    pub positive: Option<PyNumberUnaryFunc>,
129    pub absolute: Option<PyNumberUnaryFunc>,
130    pub boolean: Option<PyNumberUnaryFunc<bool>>,
131    pub invert: Option<PyNumberUnaryFunc>,
132    pub lshift: Option<PyNumberBinaryFunc>,
133    pub rshift: Option<PyNumberBinaryFunc>,
134    pub and: Option<PyNumberBinaryFunc>,
135    pub xor: Option<PyNumberBinaryFunc>,
136    pub or: Option<PyNumberBinaryFunc>,
137    pub int: Option<PyNumberUnaryFunc>,
138    pub float: Option<PyNumberUnaryFunc>,
139
140    pub inplace_add: Option<PyNumberBinaryFunc>,
141    pub inplace_subtract: Option<PyNumberBinaryFunc>,
142    pub inplace_multiply: Option<PyNumberBinaryFunc>,
143    pub inplace_remainder: Option<PyNumberBinaryFunc>,
144    pub inplace_power: Option<PyNumberTernaryFunc>,
145    pub inplace_lshift: Option<PyNumberBinaryFunc>,
146    pub inplace_rshift: Option<PyNumberBinaryFunc>,
147    pub inplace_and: Option<PyNumberBinaryFunc>,
148    pub inplace_xor: Option<PyNumberBinaryFunc>,
149    pub inplace_or: Option<PyNumberBinaryFunc>,
150
151    pub floor_divide: Option<PyNumberBinaryFunc>,
152    pub true_divide: Option<PyNumberBinaryFunc>,
153    pub inplace_floor_divide: Option<PyNumberBinaryFunc>,
154    pub inplace_true_divide: Option<PyNumberBinaryFunc>,
155
156    pub index: Option<PyNumberUnaryFunc>,
157
158    pub matrix_multiply: Option<PyNumberBinaryFunc>,
159    pub inplace_matrix_multiply: Option<PyNumberBinaryFunc>,
160}
161
162impl PyNumberMethods {
163    /// this is NOT a global variable
164    pub const NOT_IMPLEMENTED: PyNumberMethods = PyNumberMethods {
165        add: None,
166        subtract: None,
167        multiply: None,
168        remainder: None,
169        divmod: None,
170        power: None,
171        negative: None,
172        positive: None,
173        absolute: None,
174        boolean: None,
175        invert: None,
176        lshift: None,
177        rshift: None,
178        and: None,
179        xor: None,
180        or: None,
181        int: None,
182        float: None,
183        inplace_add: None,
184        inplace_subtract: None,
185        inplace_multiply: None,
186        inplace_remainder: None,
187        inplace_power: None,
188        inplace_lshift: None,
189        inplace_rshift: None,
190        inplace_and: None,
191        inplace_xor: None,
192        inplace_or: None,
193        floor_divide: None,
194        true_divide: None,
195        inplace_floor_divide: None,
196        inplace_true_divide: None,
197        index: None,
198        matrix_multiply: None,
199        inplace_matrix_multiply: None,
200    };
201
202    pub fn not_implemented() -> &'static PyNumberMethods {
203        static GLOBAL_NOT_IMPLEMENTED: PyNumberMethods = PyNumberMethods::NOT_IMPLEMENTED;
204        &GLOBAL_NOT_IMPLEMENTED
205    }
206}
207
208#[derive(Copy, Clone)]
209pub enum PyNumberBinaryOp {
210    Add,
211    Subtract,
212    Multiply,
213    Remainder,
214    Divmod,
215    Lshift,
216    Rshift,
217    And,
218    Xor,
219    Or,
220    InplaceAdd,
221    InplaceSubtract,
222    InplaceMultiply,
223    InplaceRemainder,
224    InplaceLshift,
225    InplaceRshift,
226    InplaceAnd,
227    InplaceXor,
228    InplaceOr,
229    FloorDivide,
230    TrueDivide,
231    InplaceFloorDivide,
232    InplaceTrueDivide,
233    MatrixMultiply,
234    InplaceMatrixMultiply,
235}
236
237#[derive(Copy, Clone)]
238pub enum PyNumberTernaryOp {
239    Power,
240    InplacePower,
241}
242
243#[derive(Default)]
244pub struct PyNumberSlots {
245    pub add: AtomicCell<Option<PyNumberBinaryFunc>>,
246    pub subtract: AtomicCell<Option<PyNumberBinaryFunc>>,
247    pub multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
248    pub remainder: AtomicCell<Option<PyNumberBinaryFunc>>,
249    pub divmod: AtomicCell<Option<PyNumberBinaryFunc>>,
250    pub power: AtomicCell<Option<PyNumberTernaryFunc>>,
251    pub negative: AtomicCell<Option<PyNumberUnaryFunc>>,
252    pub positive: AtomicCell<Option<PyNumberUnaryFunc>>,
253    pub absolute: AtomicCell<Option<PyNumberUnaryFunc>>,
254    pub boolean: AtomicCell<Option<PyNumberUnaryFunc<bool>>>,
255    pub invert: AtomicCell<Option<PyNumberUnaryFunc>>,
256    pub lshift: AtomicCell<Option<PyNumberBinaryFunc>>,
257    pub rshift: AtomicCell<Option<PyNumberBinaryFunc>>,
258    pub and: AtomicCell<Option<PyNumberBinaryFunc>>,
259    pub xor: AtomicCell<Option<PyNumberBinaryFunc>>,
260    pub or: AtomicCell<Option<PyNumberBinaryFunc>>,
261    pub int: AtomicCell<Option<PyNumberUnaryFunc>>,
262    pub float: AtomicCell<Option<PyNumberUnaryFunc>>,
263
264    pub right_add: AtomicCell<Option<PyNumberBinaryFunc>>,
265    pub right_subtract: AtomicCell<Option<PyNumberBinaryFunc>>,
266    pub right_multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
267    pub right_remainder: AtomicCell<Option<PyNumberBinaryFunc>>,
268    pub right_divmod: AtomicCell<Option<PyNumberBinaryFunc>>,
269    pub right_power: AtomicCell<Option<PyNumberTernaryFunc>>,
270    pub right_lshift: AtomicCell<Option<PyNumberBinaryFunc>>,
271    pub right_rshift: AtomicCell<Option<PyNumberBinaryFunc>>,
272    pub right_and: AtomicCell<Option<PyNumberBinaryFunc>>,
273    pub right_xor: AtomicCell<Option<PyNumberBinaryFunc>>,
274    pub right_or: AtomicCell<Option<PyNumberBinaryFunc>>,
275
276    pub inplace_add: AtomicCell<Option<PyNumberBinaryFunc>>,
277    pub inplace_subtract: AtomicCell<Option<PyNumberBinaryFunc>>,
278    pub inplace_multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
279    pub inplace_remainder: AtomicCell<Option<PyNumberBinaryFunc>>,
280    pub inplace_power: AtomicCell<Option<PyNumberTernaryFunc>>,
281    pub inplace_lshift: AtomicCell<Option<PyNumberBinaryFunc>>,
282    pub inplace_rshift: AtomicCell<Option<PyNumberBinaryFunc>>,
283    pub inplace_and: AtomicCell<Option<PyNumberBinaryFunc>>,
284    pub inplace_xor: AtomicCell<Option<PyNumberBinaryFunc>>,
285    pub inplace_or: AtomicCell<Option<PyNumberBinaryFunc>>,
286
287    pub floor_divide: AtomicCell<Option<PyNumberBinaryFunc>>,
288    pub true_divide: AtomicCell<Option<PyNumberBinaryFunc>>,
289    pub right_floor_divide: AtomicCell<Option<PyNumberBinaryFunc>>,
290    pub right_true_divide: AtomicCell<Option<PyNumberBinaryFunc>>,
291    pub inplace_floor_divide: AtomicCell<Option<PyNumberBinaryFunc>>,
292    pub inplace_true_divide: AtomicCell<Option<PyNumberBinaryFunc>>,
293
294    pub index: AtomicCell<Option<PyNumberUnaryFunc>>,
295
296    pub matrix_multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
297    pub right_matrix_multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
298    pub inplace_matrix_multiply: AtomicCell<Option<PyNumberBinaryFunc>>,
299}
300
301impl From<&PyNumberMethods> for PyNumberSlots {
302    fn from(value: &PyNumberMethods) -> Self {
303        // right_* functions will use the same left function as PyNumberMethods
304        // allows both f(self, other) and f(other, self)
305        Self {
306            add: AtomicCell::new(value.add),
307            subtract: AtomicCell::new(value.subtract),
308            multiply: AtomicCell::new(value.multiply),
309            remainder: AtomicCell::new(value.remainder),
310            divmod: AtomicCell::new(value.divmod),
311            power: AtomicCell::new(value.power),
312            negative: AtomicCell::new(value.negative),
313            positive: AtomicCell::new(value.positive),
314            absolute: AtomicCell::new(value.absolute),
315            boolean: AtomicCell::new(value.boolean),
316            invert: AtomicCell::new(value.invert),
317            lshift: AtomicCell::new(value.lshift),
318            rshift: AtomicCell::new(value.rshift),
319            and: AtomicCell::new(value.and),
320            xor: AtomicCell::new(value.xor),
321            or: AtomicCell::new(value.or),
322            int: AtomicCell::new(value.int),
323            float: AtomicCell::new(value.float),
324            right_add: AtomicCell::new(value.add),
325            right_subtract: AtomicCell::new(value.subtract),
326            right_multiply: AtomicCell::new(value.multiply),
327            right_remainder: AtomicCell::new(value.remainder),
328            right_divmod: AtomicCell::new(value.divmod),
329            right_power: AtomicCell::new(value.power),
330            right_lshift: AtomicCell::new(value.lshift),
331            right_rshift: AtomicCell::new(value.rshift),
332            right_and: AtomicCell::new(value.and),
333            right_xor: AtomicCell::new(value.xor),
334            right_or: AtomicCell::new(value.or),
335            inplace_add: AtomicCell::new(value.inplace_add),
336            inplace_subtract: AtomicCell::new(value.inplace_subtract),
337            inplace_multiply: AtomicCell::new(value.inplace_multiply),
338            inplace_remainder: AtomicCell::new(value.inplace_remainder),
339            inplace_power: AtomicCell::new(value.inplace_power),
340            inplace_lshift: AtomicCell::new(value.inplace_lshift),
341            inplace_rshift: AtomicCell::new(value.inplace_rshift),
342            inplace_and: AtomicCell::new(value.inplace_and),
343            inplace_xor: AtomicCell::new(value.inplace_xor),
344            inplace_or: AtomicCell::new(value.inplace_or),
345            floor_divide: AtomicCell::new(value.floor_divide),
346            true_divide: AtomicCell::new(value.true_divide),
347            right_floor_divide: AtomicCell::new(value.floor_divide),
348            right_true_divide: AtomicCell::new(value.true_divide),
349            inplace_floor_divide: AtomicCell::new(value.inplace_floor_divide),
350            inplace_true_divide: AtomicCell::new(value.inplace_true_divide),
351            index: AtomicCell::new(value.index),
352            matrix_multiply: AtomicCell::new(value.matrix_multiply),
353            right_matrix_multiply: AtomicCell::new(value.matrix_multiply),
354            inplace_matrix_multiply: AtomicCell::new(value.inplace_matrix_multiply),
355        }
356    }
357}
358
359impl PyNumberSlots {
360    pub fn left_binary_op(&self, op_slot: PyNumberBinaryOp) -> Option<PyNumberBinaryFunc> {
361        use PyNumberBinaryOp::*;
362        match op_slot {
363            Add => self.add.load(),
364            Subtract => self.subtract.load(),
365            Multiply => self.multiply.load(),
366            Remainder => self.remainder.load(),
367            Divmod => self.divmod.load(),
368            Lshift => self.lshift.load(),
369            Rshift => self.rshift.load(),
370            And => self.and.load(),
371            Xor => self.xor.load(),
372            Or => self.or.load(),
373            InplaceAdd => self.inplace_add.load(),
374            InplaceSubtract => self.inplace_subtract.load(),
375            InplaceMultiply => self.inplace_multiply.load(),
376            InplaceRemainder => self.inplace_remainder.load(),
377            InplaceLshift => self.inplace_lshift.load(),
378            InplaceRshift => self.inplace_rshift.load(),
379            InplaceAnd => self.inplace_and.load(),
380            InplaceXor => self.inplace_xor.load(),
381            InplaceOr => self.inplace_or.load(),
382            FloorDivide => self.floor_divide.load(),
383            TrueDivide => self.true_divide.load(),
384            InplaceFloorDivide => self.inplace_floor_divide.load(),
385            InplaceTrueDivide => self.inplace_true_divide.load(),
386            MatrixMultiply => self.matrix_multiply.load(),
387            InplaceMatrixMultiply => self.inplace_matrix_multiply.load(),
388        }
389    }
390
391    pub fn right_binary_op(&self, op_slot: PyNumberBinaryOp) -> Option<PyNumberBinaryFunc> {
392        use PyNumberBinaryOp::*;
393        match op_slot {
394            Add => self.right_add.load(),
395            Subtract => self.right_subtract.load(),
396            Multiply => self.right_multiply.load(),
397            Remainder => self.right_remainder.load(),
398            Divmod => self.right_divmod.load(),
399            Lshift => self.right_lshift.load(),
400            Rshift => self.right_rshift.load(),
401            And => self.right_and.load(),
402            Xor => self.right_xor.load(),
403            Or => self.right_or.load(),
404            FloorDivide => self.right_floor_divide.load(),
405            TrueDivide => self.right_true_divide.load(),
406            MatrixMultiply => self.right_matrix_multiply.load(),
407            _ => None,
408        }
409    }
410
411    pub fn left_ternary_op(&self, op_slot: PyNumberTernaryOp) -> Option<PyNumberTernaryFunc> {
412        use PyNumberTernaryOp::*;
413        match op_slot {
414            Power => self.power.load(),
415            InplacePower => self.inplace_power.load(),
416        }
417    }
418
419    pub fn right_ternary_op(&self, op_slot: PyNumberTernaryOp) -> Option<PyNumberTernaryFunc> {
420        use PyNumberTernaryOp::*;
421        match op_slot {
422            Power => self.right_power.load(),
423            _ => None,
424        }
425    }
426}
427#[derive(Copy, Clone)]
428pub struct PyNumber<'a>(&'a PyObject);
429
430unsafe impl Traverse for PyNumber<'_> {
431    fn traverse(&self, tracer_fn: &mut TraverseFn) {
432        self.0.traverse(tracer_fn)
433    }
434}
435
436impl<'a> Deref for PyNumber<'a> {
437    type Target = PyObject;
438
439    fn deref(&self) -> &Self::Target {
440        self.0
441    }
442}
443
444impl<'a> PyNumber<'a> {
445    pub(crate) fn obj(self) -> &'a PyObject {
446        self.0
447    }
448
449    // PyNumber_Check
450    pub fn check(obj: &PyObject) -> bool {
451        let methods = &obj.class().slots.as_number;
452        methods.int.load().is_some()
453            || methods.index.load().is_some()
454            || methods.float.load().is_some()
455            || obj.payload_is::<PyComplex>()
456    }
457}
458
459impl PyNumber<'_> {
460    // PyIndex_Check
461    pub fn is_index(self) -> bool {
462        self.class().slots.as_number.index.load().is_some()
463    }
464
465    #[inline]
466    pub fn int(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
467        self.class().slots.as_number.int.load().map(|f| {
468            let ret = f(self, vm)?;
469
470            if let Some(ret) = ret.downcast_ref_if_exact::<PyInt>(vm) {
471                return Ok(ret.to_owned());
472            }
473
474            let ret_class = ret.class().to_owned();
475            if let Some(ret) = ret.downcast_ref::<PyInt>() {
476                warnings::warn(
477                    vm.ctx.exceptions.deprecation_warning,
478                    format!(
479                        "__int__ returned non-int (type {}).  \
480                    The ability to return an instance of a strict subclass of int \
481                    is deprecated, and may be removed in a future version of Python.",
482                        ret_class
483                    ),
484                    1,
485                    vm,
486                )?;
487
488                Ok(ret.to_owned())
489            } else {
490                Err(vm.new_type_error(format!(
491                    "{}.__int__ returned non-int(type {})",
492                    self.class(),
493                    ret_class
494                )))
495            }
496        })
497    }
498
499    #[inline]
500    pub fn index(self, vm: &VirtualMachine) -> Option<PyResult<PyIntRef>> {
501        self.class().slots.as_number.index.load().map(|f| {
502            let ret = f(self, vm)?;
503
504            if let Some(ret) = ret.downcast_ref_if_exact::<PyInt>(vm) {
505                return Ok(ret.to_owned());
506            }
507
508            let ret_class = ret.class().to_owned();
509            if let Some(ret) = ret.downcast_ref::<PyInt>() {
510                warnings::warn(
511                    vm.ctx.exceptions.deprecation_warning,
512                    format!(
513                        "__index__ returned non-int (type {}).  \
514                    The ability to return an instance of a strict subclass of int \
515                    is deprecated, and may be removed in a future version of Python.",
516                        ret_class
517                    ),
518                    1,
519                    vm,
520                )?;
521
522                Ok(ret.to_owned())
523            } else {
524                Err(vm.new_type_error(format!(
525                    "{}.__index__ returned non-int(type {})",
526                    self.class(),
527                    ret_class
528                )))
529            }
530        })
531    }
532
533    #[inline]
534    pub fn float(self, vm: &VirtualMachine) -> Option<PyResult<PyRef<PyFloat>>> {
535        self.class().slots.as_number.float.load().map(|f| {
536            let ret = f(self, vm)?;
537
538            if let Some(ret) = ret.downcast_ref_if_exact::<PyFloat>(vm) {
539                return Ok(ret.to_owned());
540            }
541
542            let ret_class = ret.class().to_owned();
543            if let Some(ret) = ret.downcast_ref::<PyFloat>() {
544                warnings::warn(
545                    vm.ctx.exceptions.deprecation_warning,
546                    format!(
547                        "__float__ returned non-float (type {}).  \
548                    The ability to return an instance of a strict subclass of float \
549                    is deprecated, and may be removed in a future version of Python.",
550                        ret_class
551                    ),
552                    1,
553                    vm,
554                )?;
555
556                Ok(ret.to_owned())
557            } else {
558                Err(vm.new_type_error(format!(
559                    "{}.__float__ returned non-float(type {})",
560                    self.class(),
561                    ret_class
562                )))
563            }
564        })
565    }
566}
rustpython_vm/protocol/number.rs

rustpython_vm/protocol/
number.rs
