diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index ead224c6a..cce00c8f5 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1169,8 +1169,6 @@ class TestBasicOps(unittest.TestCase): p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped self.assertRaises(StopIteration, next, p) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_repeat(self): self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a']) self.assertEqual(lzip(range(3),repeat('a')), @@ -1197,8 +1195,6 @@ class TestBasicOps(unittest.TestCase): for proto in range(pickle.HIGHEST_PROTOCOL + 1): self.pickletest(proto, repeat(object='a', times=10)) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_repeat_with_negative_times(self): self.assertEqual(repr(repeat('a', -1)), "repeat('a', 0)") self.assertEqual(repr(repeat('a', -2)), "repeat('a', 0)") @@ -2181,15 +2177,11 @@ class TestVariousIteratorArgs(unittest.TestCase): class LengthTransparency(unittest.TestCase): - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_repeat(self): self.assertEqual(operator.length_hint(repeat(None, 50)), 50) self.assertEqual(operator.length_hint(repeat(None, 0)), 0) self.assertEqual(operator.length_hint(repeat(None), 12), 12) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_repeat_with_negative_times(self): self.assertEqual(operator.length_hint(repeat(None, -1)), 0) self.assertEqual(operator.length_hint(repeat(None, -2)), 0) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index f32381454..54938b3b2 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -260,7 +260,7 @@ mod decl { #[derive(Debug)] struct PyItertoolsRepeat { object: PyObjectRef, - times: Option>, + times: Option>, } impl PyValue for PyItertoolsRepeat { @@ -269,40 +269,81 @@ mod decl { } } - #[pyimpl(with(PyIter))] + #[derive(FromArgs)] + struct PyRepeatNewArgs { + #[pyarg(any)] + object: PyObjectRef, + #[pyarg(any, optional)] + times: OptionalArg, + } + + #[pyimpl(with(PyIter), flags(BASETYPE))] impl PyItertoolsRepeat { #[pyslot] fn tp_new( cls: PyTypeRef, - object: PyObjectRef, - times: OptionalArg, + PyRepeatNewArgs { object, times }: PyRepeatNewArgs, vm: &VirtualMachine, ) -> PyResult> { - let times = times - .into_option() - .map(|int| PyRwLock::new(int.as_bigint().clone())); - + let times = match times.into_option() { + Some(int) => { + let val = int.as_bigint(); + if *val > BigInt::from(isize::MAX) { + return Err(vm.new_overflow_error("Cannot fit in isize.".to_owned())); + } + // times always >= 0. + Some(PyRwLock::new(val.to_usize().unwrap_or(0))) + } + None => None, + }; PyItertoolsRepeat { object, times }.into_ref_with_type(vm, cls) } #[pymethod(name = "__length_hint__")] - fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef { + fn length_hint(&self, vm: &VirtualMachine) -> PyResult { match self.times { - Some(ref times) => vm.ctx.new_int(times.read().clone()), - None => vm.ctx.new_int(0), + Some(ref times) => Ok(vm.ctx.new_int(*times.read())), + // Return TypeError, length_hint picks this up and returns the default. + None => Err(vm.new_type_error("length of unsized object.".to_owned())), } } + + #[pymethod(magic)] + fn reduce(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + let cls = zelf.clone_class().into_pyobject(vm); + Ok(match zelf.times { + Some(ref times) => vm.ctx.new_tuple(vec![ + cls, + vm.ctx.new_tuple(vec![ + zelf.object.clone(), + vm.ctx.new_int(*times.read()).into_pyobject(vm), + ]), + ]), + None => vm + .ctx + .new_tuple(vec![cls, vm.ctx.new_tuple(vec![zelf.object.clone()])]), + }) + } + + #[pymethod(magic)] + fn repr(&self, vm: &VirtualMachine) -> PyResult { + let mut fmt = format!("{}", vm.to_repr(&self.object)?); + if let Some(ref times) = self.times { + fmt.push_str(&format!(", {}", times.read())); + } + Ok(format!("repeat({})", fmt)) + } } + impl PyIter for PyItertoolsRepeat { fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { if let Some(ref times) = zelf.times { let mut times = times.write(); - if !times.is_positive() { + if *times == 0 { return Err(vm.new_stop_iteration()); } *times -= 1; } - Ok(zelf.object.clone()) } }