Merge pull request #2676 from PyO3/exact-iter

Add more implementations of ExactSizeIterator when iterating built-in Python data structures.

Author: davidhewitt

newsfragments/2676.added.md

@@ -0,0 +1 @@
+Implemented `ExactSizeIterator` for `PyListIterator`, `PyDictIterator`, `PySetIterator` and `PyFrozenSetIterator`

src/types/dict.rs

@@ -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>;

src/types/frozenset.rs

@@ -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)
         }
     }
 }

src/types/list.rs

@@ -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)
     }
 }
 

src/types/set.rs

@@ -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)
         }
     }
 }

src/types/tuple.rs

@@ -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)
     }
 }