diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py index a2d2e3bb3..2f38550ff 100644 --- a/Lib/test/test_builtin.py +++ b/Lib/test/test_builtin.py @@ -1365,7 +1365,6 @@ class BuiltinTest(ComplexesAreIdenticalMixin, unittest.TestCase): # strict map tests based on strict zip tests - @unittest.expectedFailure # TODO: RUSTPYTHON; TypeError: Unexpected keyword argument strict def test_map_pickle_strict(self): a = (1, 2, 3) b = (4, 5, 6) @@ -1374,7 +1373,6 @@ class BuiltinTest(ComplexesAreIdenticalMixin, unittest.TestCase): m1 = map(pack, a, b, strict=True) self.check_iter_pickle(m1, t, proto) - @unittest.expectedFailure # TODO: RUSTPYTHON; TypeError: Unexpected keyword argument strict def test_map_pickle_strict_fail(self): a = (1, 2, 3) b = (4, 5, 6, 7) @@ -1385,7 +1383,6 @@ class BuiltinTest(ComplexesAreIdenticalMixin, unittest.TestCase): self.assertEqual(self.iter_error(m1, ValueError), t) self.assertEqual(self.iter_error(m2, ValueError), t) - @unittest.expectedFailure # TODO: RUSTPYTHON; TypeError: Unexpected keyword argument strict def test_map_strict(self): self.assertEqual(tuple(map(pack, (1, 2, 3), 'abc', strict=True)), ((1, 'a'), (2, 'b'), (3, 'c'))) @@ -1412,7 +1409,6 @@ class BuiltinTest(ComplexesAreIdenticalMixin, unittest.TestCase): self.assertRaises(ValueError, tuple, map(pack, 'a', t2, t3, strict=True)) - @unittest.expectedFailure # TODO: RUSTPYTHON; TypeError: Unexpected keyword argument strict def test_map_strict_iterators(self): x = iter(range(5)) y = [0] @@ -1422,7 +1418,6 @@ class BuiltinTest(ComplexesAreIdenticalMixin, unittest.TestCase): self.assertEqual(next(x), 2) self.assertEqual(next(z), 1) - @unittest.expectedFailure # TODO: RUSTPYTHON; TypeError: Unexpected keyword argument strict def test_map_strict_error_handling(self): class Error(Exception): @@ -1456,7 +1451,6 @@ class BuiltinTest(ComplexesAreIdenticalMixin, unittest.TestCase): l8 = self.iter_error(map(pack, Iter(3), "AB", strict=True), ValueError) self.assertEqual(l8, [(2, "A"), (1, "B")]) - @unittest.expectedFailure # TODO: RUSTPYTHON; TypeError: Unexpected keyword argument strict def test_map_strict_error_handling_stopiteration(self): class Iter: diff --git a/crates/vm/src/builtins/map.rs b/crates/vm/src/builtins/map.rs index d9a0e83b7..663c1bc94 100644 --- a/crates/vm/src/builtins/map.rs +++ b/crates/vm/src/builtins/map.rs @@ -1,19 +1,21 @@ -use super::{PyType, PyTypeRef}; +use super::PyType; use crate::{ - Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine, + AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine, builtins::PyTupleRef, class::PyClassImpl, - function::PosArgs, + function::{ArgIntoBool, OptionalArg, PosArgs}, protocol::{PyIter, PyIterReturn}, - raise_if_stop, types::{Constructor, IterNext, Iterable, SelfIter}, }; +use rustpython_common::atomic::{self, PyAtomic, Radium}; #[pyclass(module = false, name = "map", traverse)] #[derive(Debug)] pub struct PyMap { mapper: PyObjectRef, iterators: Vec, + #[pytraverse(skip)] + strict: PyAtomic, } impl PyPayload for PyMap { @@ -23,16 +25,27 @@ impl PyPayload for PyMap { } } +#[derive(FromArgs)] +pub struct PyMapNewArgs { + #[pyarg(named, optional)] + strict: OptionalArg, +} + impl Constructor for PyMap { - type Args = (PyObjectRef, PosArgs); + type Args = (PyObjectRef, PosArgs, PyMapNewArgs); fn py_new( _cls: &Py, - (mapper, iterators): Self::Args, + (mapper, iterators, args): Self::Args, _vm: &VirtualMachine, ) -> PyResult { let iterators = iterators.into_vec(); - Ok(Self { mapper, iterators }) + let strict = Radium::new(args.strict.unwrap_or(false)); + Ok(Self { + mapper, + iterators, + strict, + }) } } @@ -48,10 +61,24 @@ impl PyMap { } #[pymethod] - fn __reduce__(&self, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef) { - let mut vec = vec![self.mapper.clone()]; - vec.extend(self.iterators.iter().map(|o| o.clone().into())); - (vm.ctx.types.map_type.to_owned(), vm.new_tuple(vec)) + fn __reduce__(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + let cls = zelf.class().to_owned(); + let mut vec = vec![zelf.mapper.clone()]; + vec.extend(zelf.iterators.iter().map(|o| o.clone().into())); + let tuple_args = vm.ctx.new_tuple(vec); + Ok(if zelf.strict.load(atomic::Ordering::Acquire) { + vm.new_tuple((cls, tuple_args, true)) + } else { + vm.new_tuple((cls, tuple_args)) + }) + } + + #[pymethod] + fn __setstate__(zelf: PyRef, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + if let Ok(obj) = ArgIntoBool::try_from_object(vm, state) { + zelf.strict.store(obj.into(), atomic::Ordering::Release); + } + Ok(()) } } @@ -60,8 +87,35 @@ impl SelfIter for PyMap {} impl IterNext for PyMap { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { let mut next_objs = Vec::new(); - for iterator in &zelf.iterators { - let item = raise_if_stop!(iterator.next(vm)?); + for (idx, iterator) in zelf.iterators.iter().enumerate() { + let item = match iterator.next(vm)? { + PyIterReturn::Return(obj) => obj, + 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!( + "map() argument {} is shorter than argument{}{}", + idx + 1, + plural, + idx, + ))); + } + for (idx, iterator) in zelf.iterators[1..].iter().enumerate() { + if let PyIterReturn::Return(_) = iterator.next(vm)? { + let plural = if idx == 0 { " " } else { "s 1-" }; + return Err(vm.new_value_error(format!( + "map() argument {} is longer than argument{}{}", + idx + 2, + plural, + idx + 1, + ))); + } + } + } + return Ok(PyIterReturn::StopIteration(v)); + } + }; next_objs.push(item); }