rustpython_vm/builtins/
zip.rs

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

rustpython_vm/builtins/
zip.rs
