rustpython_vm/builtins/
map.rs

1use super::{PyType, PyTypeRef};
2use crate::{
3    builtins::PyTupleRef,
4    class::PyClassImpl,
5    function::PosArgs,
6    protocol::{PyIter, PyIterReturn},
7    types::{Constructor, IterNext, Iterable, SelfIter},
8    Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
9};
10
11#[pyclass(module = false, name = "map", traverse)]
12#[derive(Debug)]
13pub struct PyMap {
14    mapper: PyObjectRef,
15    iterators: Vec<PyIter>,
16}
17
18impl PyPayload for PyMap {
19    fn class(ctx: &Context) -> &'static Py<PyType> {
20        ctx.types.map_type
21    }
22}
23
24impl Constructor for PyMap {
25    type Args = (PyObjectRef, PosArgs<PyIter>);
26
27    fn py_new(cls: PyTypeRef, (mapper, iterators): Self::Args, vm: &VirtualMachine) -> PyResult {
28        let iterators = iterators.into_vec();
29        PyMap { mapper, iterators }
30            .into_ref_with_type(vm, cls)
31            .map(Into::into)
32    }
33}
34
35#[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))]
36impl PyMap {
37    #[pymethod(magic)]
38    fn length_hint(&self, vm: &VirtualMachine) -> PyResult<usize> {
39        self.iterators.iter().try_fold(0, |prev, cur| {
40            let cur = cur.as_ref().to_owned().length_hint(0, vm)?;
41            let max = std::cmp::max(prev, cur);
42            Ok(max)
43        })
44    }
45
46    #[pymethod(magic)]
47    fn reduce(&self, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef) {
48        let mut vec = vec![self.mapper.clone()];
49        vec.extend(self.iterators.iter().map(|o| o.clone().into()));
50        (vm.ctx.types.map_type.to_owned(), vm.new_tuple(vec))
51    }
52}
53
54impl SelfIter for PyMap {}
55impl IterNext for PyMap {
56    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
57        let mut next_objs = Vec::new();
58        for iterator in &zelf.iterators {
59            let item = match iterator.next(vm)? {
60                PyIterReturn::Return(obj) => obj,
61                PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
62            };
63            next_objs.push(item);
64        }
65
66        // the mapper itself can raise StopIteration which does stop the map iteration
67        PyIterReturn::from_pyresult(zelf.mapper.call(next_objs, vm), vm)
68    }
69}
70
71pub fn init(context: &Context) {
72    PyMap::extend_class(context, context.types.map_type);
73}
rustpython_vm/builtins/map.rs

rustpython_vm/builtins/
map.rs
