add length_check on `zip` via option strict param
when strict option is enabled, zipped iterators' lengths are must be
same. If not, value error will be occurred
*
use super::PyTypeRef;
use crate::{
- function::PosArgs,
+ function::{OptionalArg, PosArgs},
protocol::{PyIter, PyIterReturn},
slots::{IteratorIterable, SlotConstructor, SlotIterator},
PyClassImpl, PyContext, PyRef, PyResult, PyValue, VirtualMachine,
};
+use rustpython_common::atomic::{self, PyAtomic, Radium};
#[pyclass(module = false, name = "zip")]
#[derive(Debug)]
pub struct PyZip {
iterators: Vec<PyIter>,
+ strict: PyAtomic<bool>,
}
impl PyValue for PyZip {
@@ -18,12 +20,19 @@ impl PyValue for PyZip {
}
}
+#[derive(FromArgs)]
+pub struct PyZipNewArgs {
+ #[pyarg(named, optional)]
+ strict: OptionalArg<bool>,
+}
+
impl SlotConstructor for PyZip {
- type Args = PosArgs<PyIter>;
+ type Args = (PosArgs<PyIter>, PyZipNewArgs);
- fn py_new(cls: PyTypeRef, iterators: Self::Args, vm: &VirtualMachine) -> PyResult {
+ fn py_new(cls: PyTypeRef, (iterators, args): Self::Args, vm: &VirtualMachine) -> PyResult {
let iterators = iterators.into_vec();
- PyZip { iterators }.into_pyresult_with_type(vm, cls)
+ let strict = Radium::new(args.strict.unwrap_or(false));
+ PyZip { iterators, strict }.into_pyresult_with_type(vm, cls)
}
}
@@ -37,10 +46,34 @@ impl SlotIterator for PyZip {
return Ok(PyIterReturn::StopIteration(None));
}
let mut next_objs = Vec::new();
- for iterator in zelf.iterators.iter() {
+ for (idx, iterator) in zelf.iterators.iter().enumerate() {
let item = match iterator.next(vm)? {
PyIterReturn::Return(obj) => obj,
- PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
+ PyIterReturn::StopIteration(v) => {
+ if zelf.strict.load(atomic::Ordering::Acquire) {
+ if idx > 0 {
+ let plural = if idx == 1 { " " } else { "s 1-" };
+ return Err(vm.new_value_error(format!(
+ "zip() argument {} is shorter than argument{}{}",
+ idx + 1,
+ plural,
+ idx
+ )));
+ }
+ for (idx, iterator) in zelf.iterators[1..].iter().enumerate() {
+ if let PyIterReturn::Return(_obj) = iterator.next(vm)? {
+ let plural = if idx == 0 { " " } else { "s 1-" };
+ return Err(vm.new_value_error(format!(
+ "zip() argument {} is longer than argument{}{}",
+ idx + 2,
+ plural,
+ idx + 1
+ )));
+ }
+ }
+ }
+ return Ok(PyIterReturn::StopIteration(v));
+ }
};
next_objs.push(item);
}
