Skip to content

Commit

Permalink
Merge pull request #2676 from PyO3/exact-iter
Browse files Browse the repository at this point in the history
Add more implementations of ExactSizeIterator when iterating built-in Python data structures.
  • Loading branch information
davidhewitt authored Oct 13, 2022
2 parents 61fd70c + cd361dc commit 125af9b
Showing 6 changed files with 36 additions and 23 deletions.
1 change: 1 addition & 0 deletions newsfragments/2676.added.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Implemented `ExactSizeIterator` for `PyListIterator`, `PyDictIterator`, `PySetIterator` and `PyFrozenSetIterator`
8 changes: 7 additions & 1 deletion src/types/dict.rs
Original file line number Diff line number Diff line change
@@ -310,11 +310,17 @@ impl<'py> Iterator for PyDictIterator<'py> {

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.len as usize;
let len = self.len();
(len, Some(len))
}
}

impl<'py> ExactSizeIterator for PyDictIterator<'py> {
fn len(&self) -> usize {
self.len as usize
}
}

impl<'a> std::iter::IntoIterator for &'a PyDict {
type Item = (&'a PyAny, &'a PyAny);
type IntoIter = PyDictIterator<'a>;
15 changes: 9 additions & 6 deletions src/types/frozenset.rs
Original file line number Diff line number Diff line change
@@ -119,7 +119,7 @@ mod impl_ {

/// PyO3 implementation of an iterator for a Python `frozenset` object.
pub struct PyFrozenSetIterator<'py> {
set: &'py PyAny,
set: &'py PyFrozenSet,
pos: ffi::Py_ssize_t,
}

@@ -143,11 +143,14 @@ mod impl_ {

#[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)),
)
let len = self.len();
(len, Some(len))
}
}

impl<'py> ExactSizeIterator for PyFrozenSetIterator<'py> {
fn len(&self) -> usize {
self.set.len().saturating_sub(self.pos as usize)
}
}
}
12 changes: 7 additions & 5 deletions src/types/list.rs
Original file line number Diff line number Diff line change
@@ -322,12 +322,14 @@ impl<'a> Iterator for PyListIterator<'a> {

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.list.len();
let len = self.len();
(len, Some(len))
}
}

(
len.saturating_sub(self.index),
Some(len.saturating_sub(self.index)),
)
impl<'a> ExactSizeIterator for PyListIterator<'a> {
fn len(&self) -> usize {
self.list.len().saturating_sub(self.index)
}
}

15 changes: 9 additions & 6 deletions src/types/set.rs
Original file line number Diff line number Diff line change
@@ -169,7 +169,7 @@ mod impl_ {

/// PyO3 implementation of an iterator for a Python `set` object.
pub struct PySetIterator<'py> {
set: &'py super::PyAny,
set: &'py super::PySet,
pos: ffi::Py_ssize_t,
used: ffi::Py_ssize_t,
}
@@ -219,11 +219,14 @@ mod impl_ {

#[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)),
)
let len = self.len();
(len, Some(len))
}
}

impl<'py> ExactSizeIterator for PySetIterator<'py> {
fn len(&self) -> usize {
self.set.len().saturating_sub(self.pos as usize)
}
}
}
8 changes: 3 additions & 5 deletions src/types/tuple.rs
Original file line number Diff line number Diff line change
@@ -242,16 +242,14 @@ impl<'a> Iterator for PyTupleIterator<'a> {

#[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)),
)
let len = self.len();
(len, Some(len))
}
}

impl<'a> ExactSizeIterator for PyTupleIterator<'a> {
fn len(&self) -> usize {
self.length - self.index
self.length.saturating_sub(self.index)
}
}

0 comments on commit 125af9b

Please sign in to comment.