forked from RustPython/RustPython
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathmap.rs
More file actions
64 lines (56 loc) · 1.87 KB
/
Copy pathmap.rs
File metadata and controls
64 lines (56 loc) · 1.87 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
use super::PyTypeRef;
use crate::{
function::PosArgs,
iterator,
protocol::PyIter,
slots::{IteratorIterable, SlotConstructor, SlotIterator},
PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine,
};
/// map(func, *iterables) --> map object
///
/// Make an iterator that computes the function using arguments from
/// each of the iterables. Stops when the shortest iterable is exhausted.
#[pyclass(module = false, name = "map")]
#[derive(Debug)]
pub struct PyMap {
mapper: PyObjectRef,
iterators: Vec<PyIter>,
}
impl PyValue for PyMap {
fn class(vm: &VirtualMachine) -> &PyTypeRef {
&vm.ctx.types.map_type
}
}
impl SlotConstructor for PyMap {
type Args = (PyObjectRef, PosArgs<PyIter>);
fn py_new(cls: PyTypeRef, (mapper, iterators): Self::Args, vm: &VirtualMachine) -> PyResult {
let iterators = iterators.into_vec();
PyMap { mapper, iterators }.into_pyresult_with_type(vm, cls)
}
}
#[pyimpl(with(SlotIterator, SlotConstructor), flags(BASETYPE))]
impl PyMap {
#[pymethod(magic)]
fn length_hint(&self, vm: &VirtualMachine) -> PyResult<usize> {
self.iterators.iter().try_fold(0, |prev, cur| {
let cur = iterator::length_hint(vm, cur.as_object().clone())?.unwrap_or(0);
let max = std::cmp::max(prev, cur);
Ok(max)
})
}
}
impl IteratorIterable for PyMap {}
impl SlotIterator for PyMap {
fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
let next_objs = zelf
.iterators
.iter()
.map(|iterator| iterator.next(vm))
.collect::<Result<Vec<_>, _>>()?;
// the mapper itself can raise StopIteration which does stop the map iteration
vm.invoke(&zelf.mapper, next_objs)
}
}
pub fn init(context: &PyContext) {
PyMap::extend_class(context, &context.types.map_type);
}