Optimize Iterator::count() for PyDict, PyList, PyTuple & PySet

newsfragments/4878.added.md

@@ -1 +1,2 @@
-Optimizes `last` for `BoundListIterator`, `BoundTupleIterator` and `BorrowedTupleIterator`
+- Optimizes `last` for `BoundListIterator`, `BoundTupleIterator` and `BorrowedTupleIterator`
+- Optimizes `Iterator::count()` for `PyDict`, `PyList`, `PyTuple` & `PySet`

src/types/dict.rs

@@ -534,6 +534,14 @@ impl<'py> Iterator for BoundDictIterator<'py> {
         (len, Some(len))
     }
 
+    #[inline]
+    fn count(self) -> usize
+    where
+        Self: Sized,
+    {
+        self.len()
+    }
+
     #[inline]
     #[cfg(Py_GIL_DISABLED)]
     fn fold<B, F>(mut self, init: B, mut f: F) -> B
@@ -736,6 +744,14 @@ mod borrowed_iter {
             let len = self.len();
             (len, Some(len))
         }
+
+        #[inline]
+        fn count(self) -> usize
+        where
+            Self: Sized,
+        {
+            self.len()
+        }
     }
 
     impl ExactSizeIterator for BorrowedDictIter<'_, '_> {
@@ -1657,4 +1673,12 @@ mod tests {
                 .is_err());
         });
     }
+
+    #[test]
+    fn test_iter_count() {
+        Python::with_gil(|py| {
+            let dict = [(1, 1), (2, 2), (3, 3)].into_py_dict(py).unwrap();
+            assert_eq!(dict.iter().count(), 3);
+        })
+    }
 }

src/types/frozenset.rs

@@ -237,6 +237,14 @@ impl<'py> Iterator for BoundFrozenSetIterator<'py> {
     fn size_hint(&self) -> (usize, Option<usize>) {
         (self.remaining, Some(self.remaining))
     }
+
+    #[inline]
+    fn count(self) -> usize
+    where
+        Self: Sized,
+    {
+        self.len()
+    }
 }
 
 impl ExactSizeIterator for BoundFrozenSetIterator<'_> {
@@ -358,4 +366,12 @@ mod tests {
             assert!(!set.contains(3).unwrap());
         });
     }
+
+    #[test]
+    fn test_iter_count() {
+        Python::with_gil(|py| {
+            let set = PyFrozenSet::new(py, vec![1, 2, 3]).unwrap();
+            assert_eq!(set.iter().count(), 3);
+        })
+    }
 }

src/types/list.rs

@@ -695,6 +695,14 @@ impl<'py> Iterator for BoundListIterator<'py> {
         (len, Some(len))
     }
 
+    #[inline]
+    fn count(self) -> usize
+    where
+        Self: Sized,
+    {
+        self.len()
+    }
+
     #[inline]
     fn last(mut self) -> Option<Self::Item>
     where
@@ -1795,4 +1803,12 @@ mod tests {
             assert_eq!(last.unwrap().extract::<i32>().unwrap(), 3);
         })
     }
+
+    #[test]
+    fn test_iter_count() {
+        Python::with_gil(|py| {
+            let list = PyList::new(py, vec![1, 2, 3]).unwrap();
+            assert_eq!(list.iter().count(), 3);
+        })
+    }
 }

src/types/set.rs

@@ -267,6 +267,14 @@ impl<'py> Iterator for BoundSetIterator<'py> {
     fn size_hint(&self) -> (usize, Option<usize>) {
         (self.remaining, Some(self.remaining))
     }
+
+    #[inline]
+    fn count(self) -> usize
+    where
+        Self: Sized,
+    {
+        self.len()
+    }
 }
 
 impl ExactSizeIterator for BoundSetIterator<'_> {
@@ -479,4 +487,12 @@ mod tests {
             assert_eq!(iter.size_hint(), (0, Some(0)));
         });
     }
+
+    #[test]
+    fn test_iter_count() {
+        Python::with_gil(|py| {
+            let set = PySet::new(py, vec![1, 2, 3]).unwrap();
+            assert_eq!(set.iter().count(), 3);
+        })
+    }
 }

src/types/tuple.rs

@@ -377,6 +377,14 @@ impl<'py> Iterator for BoundTupleIterator<'py> {
         (len, Some(len))
     }
 
+    #[inline]
+    fn count(self) -> usize
+    where
+        Self: Sized,
+    {
+        self.len()
+    }
+
     #[inline]
     fn last(mut self) -> Option<Self::Item>
     where
@@ -557,6 +565,14 @@ impl<'a, 'py> Iterator for BorrowedTupleIterator<'a, 'py> {
         (len, Some(len))
     }
 
+    #[inline]
+    fn count(self) -> usize
+    where
+        Self: Sized,
+    {
+        self.len()
+    }
+
     #[inline]
     fn last(mut self) -> Option<Self::Item>
     where
@@ -1753,4 +1769,12 @@ mod tests {
             assert_eq!(last.unwrap().extract::<i32>().unwrap(), 3);
         })
     }
+
+    #[test]
+    fn test_iter_count() {
+        Python::with_gil(|py| {
+            let tuple = PyTuple::new(py, vec![1, 2, 3]).unwrap();
+            assert_eq!(tuple.iter().count(), 3);
+        })
+    }
 }