rustpython_vm/builtins/
filter.rs

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

rustpython_vm/builtins/
filter.rs
