rustpython_vm/builtins/
zip.rs

1use super::{PyType, PyTypeRef};
2use crate::{
3    builtins::PyTupleRef,
4    class::PyClassImpl,
5    function::{ArgIntoBool, OptionalArg, PosArgs},
6    protocol::{PyIter, PyIterReturn},
7    types::{Constructor, IterNext, Iterable, SelfIter},
8    AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
9};
10use rustpython_common::atomic::{self, PyAtomic, Radium};
11
12#[pyclass(module = false, name = "zip", traverse)]
13#[derive(Debug)]
14pub struct PyZip {
15    iterators: Vec<PyIter>,
16    #[pytraverse(skip)]
17    strict: PyAtomic<bool>,
18}
19
20impl PyPayload for PyZip {
21    fn class(ctx: &Context) -> &'static Py<PyType> {
22        ctx.types.zip_type
23    }
24}
25
26#[derive(FromArgs)]
27pub struct PyZipNewArgs {
28    #[pyarg(named, optional)]
29    strict: OptionalArg<bool>,
30}
31
32impl Constructor for PyZip {
33    type Args = (PosArgs<PyIter>, PyZipNewArgs);
34
35    fn py_new(cls: PyTypeRef, (iterators, args): Self::Args, vm: &VirtualMachine) -> PyResult {
36        let iterators = iterators.into_vec();
37        let strict = Radium::new(args.strict.unwrap_or(false));
38        PyZip { iterators, strict }
39            .into_ref_with_type(vm, cls)
40            .map(Into::into)
41    }
42}
43
44#[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))]
45impl PyZip {
46    #[pymethod(magic)]
47    fn reduce(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
48        let cls = zelf.class().to_owned();
49        let iterators = zelf
50            .iterators
51            .iter()
52            .map(|obj| obj.clone().into())
53            .collect::<Vec<_>>();
54        let tuple_iter = vm.ctx.new_tuple(iterators);
55        Ok(if zelf.strict.load(atomic::Ordering::Acquire) {
56            vm.new_tuple((cls, tuple_iter, true))
57        } else {
58            vm.new_tuple((cls, tuple_iter))
59        })
60    }
61
62    #[pymethod(magic)]
63    fn setstate(zelf: PyRef<Self>, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
64        if let Ok(obj) = ArgIntoBool::try_from_object(vm, state) {
65            zelf.strict.store(obj.into(), atomic::Ordering::Release);
66        }
67        Ok(())
68    }
69}
70
71impl SelfIter for PyZip {}
72impl IterNext for PyZip {
73    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
74        if zelf.iterators.is_empty() {
75            return Ok(PyIterReturn::StopIteration(None));
76        }
77        let mut next_objs = Vec::new();
78        for (idx, iterator) in zelf.iterators.iter().enumerate() {
79            let item = match iterator.next(vm)? {
80                PyIterReturn::Return(obj) => obj,
81                PyIterReturn::StopIteration(v) => {
82                    if zelf.strict.load(atomic::Ordering::Acquire) {
83                        if idx > 0 {
84                            let plural = if idx == 1 { " " } else { "s 1-" };
85                            return Err(vm.new_value_error(format!(
86                                "zip() argument {} is shorter than argument{}{}",
87                                idx + 1,
88                                plural,
89                                idx
90                            )));
91                        }
92                        for (idx, iterator) in zelf.iterators[1..].iter().enumerate() {
93                            if let PyIterReturn::Return(_obj) = iterator.next(vm)? {
94                                let plural = if idx == 0 { " " } else { "s 1-" };
95                                return Err(vm.new_value_error(format!(
96                                    "zip() argument {} is longer than argument{}{}",
97                                    idx + 2,
98                                    plural,
99                                    idx + 1
100                                )));
101                            }
102                        }
103                    }
104                    return Ok(PyIterReturn::StopIteration(v));
105                }
106            };
107            next_objs.push(item);
108        }
109        Ok(PyIterReturn::Return(vm.ctx.new_tuple(next_objs).into()))
110    }
111}
112
113pub fn init(ctx: &Context) {
114    PyZip::extend_class(ctx, ctx.types.zip_type);
115}
rustpython_vm/builtins/zip.rs

rustpython_vm/builtins/
zip.rs
