Skip to main content

rustpython_vm/function/
number.rs

1use super::argument::OptionalArg;
2use crate::{AsObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine, builtins::PyIntRef};
3use core::ops::Deref;
4use malachite_bigint::BigInt;
5use num_complex::Complex64;
6use num_traits::PrimInt;
7
8/// A Python complex-like object.
9///
10/// `ArgIntoComplex` implements `FromArgs` so that a built-in function can accept
11/// any object that can be transformed into a complex.
12///
13/// If the object is not a Python complex object but has a `__complex__()`
14/// method, this method will first be called to convert the object into a float.
15/// If `__complex__()` is not defined then it falls back to `__float__()`. If
16/// `__float__()` is not defined it falls back to `__index__()`.
17#[derive(Debug, PartialEq)]
18#[repr(transparent)]
19pub struct ArgIntoComplex {
20    value: Complex64,
21}
22
23impl ArgIntoComplex {
24    #[inline]
25    pub fn into_complex(self) -> Complex64 {
26        self.value
27    }
28}
29
30impl From<ArgIntoComplex> for Complex64 {
31    fn from(arg: ArgIntoComplex) -> Self {
32        arg.value
33    }
34}
35
36impl TryFromObject for ArgIntoComplex {
37    // Equivalent to PyComplex_AsCComplex
38    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
39        // We do not care if it was already a complex.
40        let (value, _) = obj.try_complex(vm)?.ok_or_else(|| {
41            vm.new_type_error(format!("must be real number, not {}", obj.class().name()))
42        })?;
43        Ok(Self { value })
44    }
45}
46
47/// A Python float-like object.
48///
49/// `ArgIntoFloat` implements `FromArgs` so that a built-in function can accept
50/// any object that can be transformed into a float.
51///
52/// If the object is not a Python floating point object but has a `__float__()`
53/// method, this method will first be called to convert the object into a float.
54/// If `__float__()` is not defined then it falls back to `__index__()`.
55#[derive(Debug, PartialEq)]
56#[repr(transparent)]
57pub struct ArgIntoFloat {
58    value: f64,
59}
60
61impl ArgIntoFloat {
62    #[inline]
63    pub fn into_float(self) -> f64 {
64        self.value
65    }
66
67    pub fn vec_into_f64(v: Vec<Self>) -> Vec<f64> {
68        // TODO: Vec::into_raw_parts once stabilized
69        let mut v = core::mem::ManuallyDrop::new(v);
70        let (p, l, c) = (v.as_mut_ptr(), v.len(), v.capacity());
71        // SAFETY: IntoPyFloat is repr(transparent) over f64
72        unsafe { Vec::from_raw_parts(p.cast(), l, c) }
73    }
74}
75
76impl From<ArgIntoFloat> for f64 {
77    fn from(arg: ArgIntoFloat) -> Self {
78        arg.value
79    }
80}
81
82impl TryFromObject for ArgIntoFloat {
83    // Equivalent to PyFloat_AsDouble.
84    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
85        let value = obj.try_float(vm)?.to_f64();
86        Ok(Self { value })
87    }
88}
89
90/// A Python bool-like object.
91///
92/// `ArgIntoBool` implements `FromArgs` so that a built-in function can accept
93/// any object that can be transformed into a boolean.
94///
95/// By default an object is considered true unless its class defines either a
96/// `__bool__()` method that returns False or a `__len__()` method that returns
97/// zero, when called with the object.
98#[derive(Debug, Default, PartialEq, Eq)]
99pub struct ArgIntoBool {
100    value: bool,
101}
102
103impl ArgIntoBool {
104    pub const TRUE: Self = Self { value: true };
105    pub const FALSE: Self = Self { value: false };
106
107    #[inline]
108    pub fn into_bool(self) -> bool {
109        self.value
110    }
111}
112
113impl From<ArgIntoBool> for bool {
114    fn from(arg: ArgIntoBool) -> Self {
115        arg.value
116    }
117}
118
119impl TryFromObject for ArgIntoBool {
120    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
121        Ok(Self {
122            value: obj.try_to_bool(vm)?,
123        })
124    }
125}
126
127// Implement ArgIndex to separate between "true" int and int generated by index
128#[derive(Debug, Traverse)]
129#[repr(transparent)]
130pub struct ArgIndex {
131    value: PyIntRef,
132}
133
134impl ArgIndex {
135    #[inline]
136    pub fn into_int_ref(self) -> PyIntRef {
137        self.value
138    }
139}
140
141impl AsRef<PyIntRef> for ArgIndex {
142    fn as_ref(&self) -> &PyIntRef {
143        &self.value
144    }
145}
146
147impl From<ArgIndex> for PyIntRef {
148    fn from(arg: ArgIndex) -> Self {
149        arg.value
150    }
151}
152
153impl TryFromObject for ArgIndex {
154    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
155        Ok(Self {
156            value: obj.try_index(vm)?,
157        })
158    }
159}
160
161#[derive(Debug, Copy, Clone)]
162#[repr(transparent)]
163pub struct ArgPrimitiveIndex<T> {
164    pub value: T,
165}
166
167impl<T> OptionalArg<ArgPrimitiveIndex<T>> {
168    pub fn into_primitive(self) -> OptionalArg<T> {
169        self.map(|x| x.value)
170    }
171}
172
173impl<T> Deref for ArgPrimitiveIndex<T> {
174    type Target = T;
175
176    fn deref(&self) -> &Self::Target {
177        &self.value
178    }
179}
180
181impl<T> TryFromObject for ArgPrimitiveIndex<T>
182where
183    T: PrimInt + for<'a> TryFrom<&'a BigInt>,
184{
185    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
186        Ok(Self {
187            value: obj.try_index(vm)?.try_to_primitive(vm)?,
188        })
189    }
190}
191
192pub type ArgSize = ArgPrimitiveIndex<isize>;
193
194impl From<ArgSize> for isize {
195    fn from(arg: ArgSize) -> Self {
196        arg.value
197    }
198}