Added `size_hint` impls for `{PyDict,PyList,PySet,PyTuple}Iterator`s

This commit is contained in:
Ohad Ravid 2021-06-27 09:45:15 +03:00
parent 11f1a1d6f0
commit fcffcdfae5
5 changed files with 104 additions and 0 deletions

View File

@ -115,6 +115,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
- Add `serde` feature which provides implementations of `Serialize` and `Deserialize` for `Py<T>`. [#1366](https://github.com/PyO3/pyo3/pull/1366)
- Add FFI definition `_PyCFunctionFastWithKeywords` on Python 3.7 and up. [#1384](https://github.com/PyO3/pyo3/pull/1384)
- Add `PyDateTime::new_with_fold()` method. [#1398](https://github.com/PyO3/pyo3/pull/1398)
- Add `size_hint` impls for `{PyDict,PyList,PySet,PyTuple}Iterator`s. [#1699](https://github.com/PyO3/pyo3/pull/1699)
### Changed

View File

@ -204,6 +204,15 @@ impl<'py> Iterator for PyDictIterator<'py> {
}
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.dict.len().unwrap_or_default();
(
len.saturating_sub(self.pos as usize),
Some(len.saturating_sub(self.pos as usize)),
)
}
}
impl<'a> std::iter::IntoIterator for &'a PyDict {
@ -698,6 +707,28 @@ mod test {
assert_eq!(32 + 42 + 123, value_sum);
}
#[test]
fn test_iter_size_hint() {
let gil = Python::acquire_gil();
let py = gil.python();
let mut v = HashMap::new();
v.insert(7, 32);
v.insert(8, 42);
v.insert(9, 123);
let ob = v.to_object(py);
let dict = <PyDict as PyTryFrom>::try_from(ob.as_ref(py)).unwrap();
let mut iter = dict.iter();
assert_eq!(iter.size_hint(), (v.len(), Some(v.len())));
iter.next();
assert_eq!(iter.size_hint(), (v.len() - 1, Some(v.len() - 1)));
// Exhust iterator.
while let Some(_) = iter.next() {}
assert_eq!(iter.size_hint(), (0, Some(0)));
}
#[test]
fn test_into_iter() {
let gil = Python::acquire_gil();

View File

@ -158,6 +158,16 @@ impl<'a> Iterator for PyListIterator<'a> {
None
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.list.len();
(
len.saturating_sub(self.index as usize),
Some(len.saturating_sub(self.index as usize)),
)
}
}
impl<'a> std::iter::IntoIterator for &'a PyList {
@ -342,6 +352,25 @@ mod test {
assert_eq!(idx, v.len());
}
#[test]
fn test_iter_size_hint() {
let gil = Python::acquire_gil();
let py = gil.python();
let v = vec![2, 3, 5, 7];
let ob = v.to_object(py);
let list = <PyList as PyTryFrom>::try_from(ob.as_ref(py)).unwrap();
let mut iter = list.iter();
assert_eq!(iter.size_hint(), (v.len(), Some(v.len())));
iter.next();
assert_eq!(iter.size_hint(), (v.len() - 1, Some(v.len() - 1)));
// Exhust iterator.
while let Some(_) = iter.next() {}
assert_eq!(iter.size_hint(), (0, Some(0)));
}
#[test]
fn test_into_iter() {
let gil = Python::acquire_gil();

View File

@ -172,6 +172,15 @@ impl<'py> Iterator for PySetIterator<'py> {
}
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.set.len().unwrap_or_default();
(
len.saturating_sub(self.pos as usize),
Some(len.saturating_sub(self.pos as usize)),
)
}
}
impl<'a> std::iter::IntoIterator for &'a PySet {
@ -505,6 +514,24 @@ mod test {
}
}
#[test]
fn test_set_iter_size_hint() {
let gil = Python::acquire_gil();
let py = gil.python();
let set = PySet::new(py, &[1]).unwrap();
let mut iter = set.iter();
if cfg!(Py_LIMITED_API) {
assert_eq!(iter.size_hint(), (0, None));
} else {
assert_eq!(iter.size_hint(), (1, Some(1)));
iter.next();
assert_eq!(iter.size_hint(), (0, Some(0)));
}
}
#[test]
fn test_frozenset_new_and_len() {
let gil = Python::acquire_gil();

View File

@ -131,6 +131,14 @@ impl<'a> Iterator for PyTupleIterator<'a> {
None
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(
self.length.saturating_sub(self.index as usize),
Some(self.length.saturating_sub(self.index as usize)),
)
}
}
impl<'a> ExactSizeIterator for PyTupleIterator<'a> {
@ -346,9 +354,17 @@ mod test {
let tuple = <PyTuple as PyTryFrom>::try_from(ob.as_ref(py)).unwrap();
assert_eq!(3, tuple.len());
let mut iter = tuple.iter();
assert_eq!(iter.size_hint(), (3, Some(3)));
assert_eq!(1, iter.next().unwrap().extract().unwrap());
assert_eq!(iter.size_hint(), (2, Some(2)));
assert_eq!(2, iter.next().unwrap().extract().unwrap());
assert_eq!(iter.size_hint(), (1, Some(1)));
assert_eq!(3, iter.next().unwrap().extract().unwrap());
assert_eq!(iter.size_hint(), (0, Some(0)));
}
#[test]