rustpython_vm/function/

protocol.rs

use super::IntoFuncArgs;
use crate::{
    AsObject, PyObject, PyObjectRef, PyPayload, PyResult, TryFromObject, VirtualMachine,
    builtins::{PyDictRef, iter::PySequenceIterator},
    convert::ToPyObject,
    object::{Traverse, TraverseFn},
    protocol::{PyIter, PyIterIter, PyMapping},
    types::GenericMethod,
};
use core::{borrow::Borrow, marker::PhantomData};

#[derive(Clone, Traverse)]
pub struct ArgCallable {
    obj: PyObjectRef,
    #[pytraverse(skip)]
    call: GenericMethod,
}

impl ArgCallable {
    #[inline(always)]
    pub fn invoke(&self, args: impl IntoFuncArgs, vm: &VirtualMachine) -> PyResult {
        let args = args.into_args(vm);
        (self.call)(&self.obj, args, vm)
    }
}

impl core::fmt::Debug for ArgCallable {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("ArgCallable")
            .field("obj", &self.obj)
            .field("call", &format!("{:08x}", self.call as usize))
            .finish()
    }
}

impl Borrow<PyObject> for ArgCallable {
    #[inline(always)]
    fn borrow(&self) -> &PyObject {
        &self.obj
    }
}

impl AsRef<PyObject> for ArgCallable {
    #[inline(always)]
    fn as_ref(&self) -> &PyObject {
        &self.obj
    }
}

impl From<ArgCallable> for PyObjectRef {
    #[inline(always)]
    fn from(value: ArgCallable) -> Self {
        value.obj
    }
}

impl TryFromObject for ArgCallable {
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        let Some(callable) = obj.to_callable() else {
            return Err(vm.new_type_error(format!(
                "'{}' object is not callable",
                obj.class().slot_name()
            )));
        };
        let call = callable.call;
        Ok(Self { obj, call })
    }
}

/// An iterable Python object.
///
/// `ArgIterable` implements `FromArgs` so that a built-in function can accept
/// an object that is required to conform to the Python iterator protocol.
///
/// ArgIterable can optionally perform type checking and conversions on iterated
/// objects using a generic type parameter that implements `TryFromObject`.
pub struct ArgIterable<T = PyObjectRef> {
    iterable: PyObjectRef,
    iter_fn: Option<crate::types::IterFunc>,
    _item: PhantomData<T>,
}

unsafe impl<T: Traverse> Traverse for ArgIterable<T> {
    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
        self.iterable.traverse(tracer_fn)
    }
}

impl<T> ArgIterable<T> {
    #[must_use]
    pub(crate) fn as_object(&self) -> &PyObject {
        &self.iterable
    }

    /// This object's iterator.
    ///
    /// This operation may fail if an exception is raised while invoking the
    /// `__iter__` method of the iterable object.
    fn get_iter(&self, vm: &VirtualMachine) -> PyResult<PyIter> {
        Ok(PyIter::new(match self.iter_fn {
            Some(f) => f(self.iterable.clone(), vm)?,
            None => PySequenceIterator::new(self.iterable.clone(), vm)?.into_pyobject(vm),
        }))
    }

    /// Returns an iterator over this sequence of objects. See [`PyIter::iter`]
    /// for why it does not ask how long the iterator is.
    ///
    /// This operation may fail if an exception is raised while invoking the
    /// `__iter__` method of the iterable object.
    pub fn iter<'a>(&self, vm: &'a VirtualMachine) -> PyResult<PyIterIter<'a, T>> {
        Ok(self.get_iter(vm)?.into_iter(vm))
    }

    /// [`Self::iter`] for a caller that fills a sized container from the
    /// iterator, the way `PySequence_Fast()` does.
    pub fn iter_sized<'a>(&self, vm: &'a VirtualMachine) -> PyResult<PyIterIter<'a, T>> {
        self.get_iter(vm)?.into_iter_sized(vm)
    }
}

impl<T> TryFromObject for ArgIterable<T>
where
    T: TryFromObject,
{
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        let cls = obj.class();
        let iter_fn = cls.slots.iter.load();
        if iter_fn.is_none() && !cls.has_attr(identifier!(vm, __getitem__)) {
            return Err(vm.new_type_error(format!("'{}' object is not iterable", cls.slot_name())));
        }
        Ok(Self {
            iterable: obj,
            iter_fn,
            _item: PhantomData,
        })
    }
}

#[derive(Debug, Clone, Traverse)]
pub struct ArgMapping {
    obj: PyObjectRef,
}

impl ArgMapping {
    #[inline]
    #[must_use]
    pub const fn new(obj: PyObjectRef) -> Self {
        Self { obj }
    }

    #[inline(always)]
    #[must_use]
    pub fn from_dict_exact(dict: PyDictRef) -> Self {
        Self { obj: dict.into() }
    }

    #[inline(always)]
    #[must_use]
    pub fn obj(&self) -> &PyObject {
        &self.obj
    }

    #[inline(always)]
    #[must_use]
    pub fn mapping(&self) -> PyMapping<'_> {
        self.obj.mapping_unchecked()
    }
}

impl Borrow<PyObject> for ArgMapping {
    #[inline(always)]
    fn borrow(&self) -> &PyObject {
        &self.obj
    }
}

impl AsRef<PyObject> for ArgMapping {
    #[inline(always)]
    fn as_ref(&self) -> &PyObject {
        &self.obj
    }
}

impl From<ArgMapping> for PyObjectRef {
    #[inline(always)]
    fn from(value: ArgMapping) -> Self {
        value.obj
    }
}

impl ToPyObject for ArgMapping {
    #[inline(always)]
    fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
        self.obj
    }
}

impl TryFromObject for ArgMapping {
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        let _mapping = obj.try_mapping(vm)?;
        Ok(Self { obj })
    }
}

// this is not strictly related to PySequence protocol.
#[derive(Clone)]
pub struct ArgSequence<T = PyObjectRef>(Vec<T>);

unsafe impl<T: Traverse> Traverse for ArgSequence<T> {
    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
        self.0.traverse(tracer_fn);
    }
}

impl<T> ArgSequence<T> {
    #[inline(always)]
    #[must_use]
    pub fn into_vec(self) -> Vec<T> {
        self.0
    }
    #[inline(always)]
    #[must_use]
    pub fn as_slice(&self) -> &[T] {
        &self.0
    }
}

impl<T> core::ops::Deref for ArgSequence<T> {
    type Target = [T];
    #[inline(always)]
    fn deref(&self) -> &[T] {
        self.as_slice()
    }
}

impl<'a, T> IntoIterator for &'a ArgSequence<T> {
    type Item = &'a T;
    type IntoIter = core::slice::Iter<'a, T>;
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}
impl<T> IntoIterator for ArgSequence<T> {
    type Item = T;
    type IntoIter = alloc::vec::IntoIter<T>;
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<T: TryFromObject> TryFromObject for ArgSequence<T> {
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        obj.try_to_value(vm).map(Self)
    }
}