diff --git a/Lib/test/test_iterlen.py b/Lib/test/test_iterlen.py index 39e72169a..5e4b7cf7e 100644 --- a/Lib/test/test_iterlen.py +++ b/Lib/test/test_iterlen.py @@ -252,8 +252,6 @@ class NoneLengthHint(object): class TestLengthHintExceptions(unittest.TestCase): - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_issue1242657(self): self.assertRaises(RuntimeError, list, BadLen()) self.assertRaises(RuntimeError, list, BadLengthHint()) diff --git a/vm/src/iterator.rs b/vm/src/iterator.rs index 2bf41144e..f99bd2fd0 100644 --- a/vm/src/iterator.rs +++ b/vm/src/iterator.rs @@ -6,7 +6,7 @@ use crate::builtins::int::{self, PyInt}; use crate::builtins::iter::PySequenceIterator; use crate::exceptions::PyBaseExceptionRef; use crate::vm::VirtualMachine; -use crate::{IdProtocol, PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol}; +use crate::{IdProtocol, PyObjectRef, PyResult, PyValue, TypeProtocol}; use num_traits::Signed; /* @@ -73,16 +73,15 @@ pub fn get_next_object( } } -/* Retrieve all elements from an iterator */ -pub fn get_all(vm: &VirtualMachine, iter_obj: &PyObjectRef) -> PyResult> { - try_map(vm, iter_obj, |obj| T::try_from_object(vm, obj)) -} - -pub fn try_map(vm: &VirtualMachine, iter_obj: &PyObjectRef, mut f: F) -> PyResult> +pub fn try_map( + vm: &VirtualMachine, + iter_obj: &PyObjectRef, + cap: usize, + mut f: F, +) -> PyResult> where F: FnMut(PyObjectRef) -> PyResult, { - let cap = length_hint(vm, iter_obj.clone())?.unwrap_or(0); // TODO: fix extend to do this check (?), see test_extend in Lib/test/list_tests.py, // https://github.com/python/cpython/blob/v3.9.0/Objects/listobject.c#L922-L928 if cap >= isize::max_value() as usize { diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs index 86894c3ed..a4446ed2c 100644 --- a/vm/src/stdlib/collections.rs +++ b/vm/src/stdlib/collections.rs @@ -204,7 +204,6 @@ mod _collections { #[pymethod] fn extend(&self, iter: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - // TODO: use length_hint here and for extendleft self.state.fetch_add(1); let max_len = self.maxlen; let mut elements: Vec = vm.extract_elements(&iter)?; diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 8ac21d4c0..3059303e2 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -1261,7 +1261,13 @@ impl VirtualMachine { .collect() } else { let iter = iterator::get_iter(self, value.clone())?; - iterator::try_map(self, &iter, |obj| func(obj)) + let cap = match iterator::length_hint(self, value.clone()) { + Err(e) if e.class().is(&self.ctx.exceptions.runtime_error) => return Err(e), + Ok(Some(value)) => value, + // Use a power of 2 as a default capacity. + _ => 4, + }; + iterator::try_map(self, &iter, cap, |obj| func(obj)) } } @@ -1303,7 +1309,15 @@ impl VirtualMachine { // TODO: put internal iterable type obj => { let iter = iterator::get_iter(self, obj.clone())?; - Ok(iterator::try_map(self, &iter, f)) + let cap = match iterator::length_hint(self, obj.clone()) { + Err(e) if e.class().is(&self.ctx.exceptions.runtime_error) => { + return Ok(Err(e)) + } + Ok(Some(value)) => value, + // Use a power of 2 as a default capacity. + _ => 4, + }; + Ok(iterator::try_map(self, &iter, cap, f)) } }) }