rustpython_vm/function/

number.rs

use super::argument::OptionalArg;
use crate::{
    AsObject, Py, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
    builtins::{PyInt, PyIntRef},
};
use core::ops::Deref;
use malachite_bigint::BigInt;
use num_complex::Complex64;
use num_traits::PrimInt;

/// A Python complex-like object.
///
/// `ArgIntoComplex` implements `FromArgs` so that a built-in function can accept
/// any object that can be transformed into a complex.
///
/// If the object is not a Python complex object but has a `__complex__()`
/// method, this method will first be called to convert the object into a float.
/// If `__complex__()` is not defined then it falls back to `__float__()`. If
/// `__float__()` is not defined it falls back to `__index__()`.
#[derive(Debug, PartialEq)]
#[repr(transparent)]
pub struct ArgIntoComplex {
    value: Complex64,
}

impl ArgIntoComplex {
    #[inline]
    #[must_use]
    pub fn into_complex(self) -> Complex64 {
        self.value
    }
}

impl From<ArgIntoComplex> for Complex64 {
    fn from(arg: ArgIntoComplex) -> Self {
        arg.value
    }
}

impl TryFromObject for ArgIntoComplex {
    // Equivalent to PyComplex_AsCComplex
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        // We do not care if it was already a complex.
        let (value, _) = obj.try_complex(vm)?.ok_or_else(|| {
            vm.new_type_error(format!("must be real number, not {}", obj.class().name()))
        })?;
        Ok(Self { value })
    }
}

/// A Python float-like object.
///
/// `ArgIntoFloat` implements `FromArgs` so that a built-in function can accept
/// any object that can be transformed into a float.
///
/// If the object is not a Python floating point object but has a `__float__()`
/// method, this method will first be called to convert the object into a float.
/// If `__float__()` is not defined then it falls back to `__index__()`.
#[derive(Debug, PartialEq)]
#[repr(transparent)]
pub struct ArgIntoFloat {
    value: f64,
}

impl ArgIntoFloat {
    #[inline]
    #[must_use]
    pub fn into_float(self) -> f64 {
        self.value
    }

    #[must_use]
    pub fn vec_into_f64(v: Vec<Self>) -> Vec<f64> {
        // TODO: Vec::into_raw_parts once stabilized
        let mut v = core::mem::ManuallyDrop::new(v);
        let (p, l, c) = (v.as_mut_ptr(), v.len(), v.capacity());
        // SAFETY: IntoPyFloat is repr(transparent) over f64
        unsafe { Vec::from_raw_parts(p.cast(), l, c) }
    }
}

impl From<ArgIntoFloat> for f64 {
    fn from(arg: ArgIntoFloat) -> Self {
        arg.value
    }
}

impl TryFromObject for ArgIntoFloat {
    // Equivalent to PyFloat_AsDouble.
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        let value = obj.try_float(vm)?.to_f64();
        Ok(Self { value })
    }
}

/// A Python bool-like object.
///
/// `ArgIntoBool` implements `FromArgs` so that a built-in function can accept
/// any object that can be transformed into a boolean.
///
/// By default an object is considered true unless its class defines either a
/// `__bool__()` method that returns False or a `__len__()` method that returns
/// zero, when called with the object.
#[derive(Debug, Default, PartialEq, Eq)]
pub struct ArgIntoBool {
    value: bool,
}

impl ArgIntoBool {
    pub const TRUE: Self = Self { value: true };
    pub const FALSE: Self = Self { value: false };

    #[inline]
    #[must_use]
    pub fn into_bool(self) -> bool {
        self.value
    }

    #[must_use]
    pub const fn py_default(&self) -> super::DefaultRepr {
        super::DefaultRepr::Bool(self.value)
    }
}

impl From<ArgIntoBool> for bool {
    fn from(arg: ArgIntoBool) -> Self {
        arg.value
    }
}

impl TryFromObject for ArgIntoBool {
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        Ok(Self {
            value: obj.try_to_bool(vm)?,
        })
    }
}

// Implement ArgIndex to separate between "true" int and int generated by index
#[derive(Debug, Traverse)]
#[repr(transparent)]
pub struct ArgIndex {
    value: PyIntRef,
}

impl ArgIndex {
    #[inline]
    #[must_use]
    pub fn into_int_ref(self) -> PyIntRef {
        self.value
    }
}

impl AsRef<Py<PyInt>> for ArgIndex {
    fn as_ref(&self) -> &Py<PyInt> {
        &self.value
    }
}

impl From<ArgIndex> for PyIntRef {
    fn from(arg: ArgIndex) -> Self {
        arg.value
    }
}

impl TryFromObject for ArgIndex {
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        Ok(Self {
            value: obj.try_index(vm)?,
        })
    }
}

/// A signed size or index argument (`Py_ssize_t`).
pub type PySsize = isize;
/// An unsigned size argument (`size_t`).
pub type PySize = usize;

/// An `int` (or subclass, including `bool`) converted to a Rust primitive.
///
/// Unlike the primitive `TryFromObject` impls, this never calls `__index__`.
#[derive(Debug, Copy, Clone)]
#[repr(transparent)]
pub struct ArgStrictInt<T> {
    pub value: T,
}

impl<T> ArgStrictInt<T> {
    #[inline]
    #[must_use]
    pub fn into_primitive(self) -> T {
        self.value
    }
}

impl<T> OptionalArg<ArgStrictInt<T>> {
    pub fn into_primitive(self) -> OptionalArg<T> {
        self.map(|x| x.value)
    }
}

impl<T> Deref for ArgStrictInt<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

impl<T> TryFromObject for ArgStrictInt<T>
where
    T: PrimInt + for<'a> TryFrom<&'a BigInt>,
{
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        // Same check the primitive conversions used before they grew `__index__`.
        let value = obj.try_value_with(|int: &Py<PyInt>| int.try_to_primitive(vm), vm)?;
        Ok(Self { value })
    }
}