always use a Python iterator for sets and frozensets

PyO3 · Feb 17, 2024 · e95050c · e95050c
1 parent c24478e
commit e95050c
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 184 deletions.
diff --git a/src/types/frozenset.rs b/src/types/frozenset.rs
@@ -1,4 +1,3 @@
-#[cfg(Py_LIMITED_API)]
 use crate::types::PyIterator;
 use crate::{
     err::{self, PyErr, PyResult},
@@ -205,6 +204,13 @@ impl<'py> Iterator for PyFrozenSetIterator<'py> {
     }
 }
 
+impl ExactSizeIterator for PyFrozenSetIterator<'_> {
+    #[inline]
+    fn len(&self) -> usize {
+        self.0.len()
+    }
+}
+
 impl<'py> IntoIterator for &'py PyFrozenSet {
     type Item = &'py PyAny;
     type IntoIter = PyFrozenSetIterator<'py>;
@@ -224,90 +230,42 @@ impl<'py> IntoIterator for Bound<'py, PyFrozenSet> {
     }
 }
 
-#[cfg(Py_LIMITED_API)]
-mod impl_ {
-    use super::*;
-
-    /// PyO3 implementation of an iterator for a Python `set` object.
-    pub struct BoundFrozenSetIterator<'p> {
-        it: Bound<'p, PyIterator>,
-    }
-
-    impl<'py> BoundFrozenSetIterator<'py> {
-        pub(super) fn new(frozenset: Bound<'py, PyFrozenSet>) -> Self {
-            Self {
-                it: PyIterator::from_bound_object(&frozenset).unwrap(),
-            }
-        }
-    }
-
-    impl<'py> Iterator for BoundFrozenSetIterator<'py> {
-        type Item = Bound<'py, super::PyAny>;
-
-        /// Advances the iterator and returns the next value.
-        #[inline]
-        fn next(&mut self) -> Option<Self::Item> {
-            self.it.next().map(Result::unwrap)
-        }
-    }
+/// PyO3 implementation of an iterator for a Python `frozenset` object.
+pub struct BoundFrozenSetIterator<'p> {
+    it: Bound<'p, PyIterator>,
+    // Remaining elements in the frozenset
+    remaining: usize,
 }
 
-#[cfg(not(Py_LIMITED_API))]
-mod impl_ {
-    use super::*;
-
-    /// PyO3 implementation of an iterator for a Python `frozenset` object.
-    pub struct BoundFrozenSetIterator<'py> {
-        set: Bound<'py, PyFrozenSet>,
-        pos: ffi::Py_ssize_t,
-    }
-
-    impl<'py> BoundFrozenSetIterator<'py> {
-        pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
-            Self { set, pos: 0 }
+impl<'py> BoundFrozenSetIterator<'py> {
+    pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
+        Self {
+            it: PyIterator::from_bound_object(&set).unwrap(),
+            remaining: set.len(),
         }
     }
+}
 
-    impl<'py> Iterator for BoundFrozenSetIterator<'py> {
-        type Item = Bound<'py, PyAny>;
-
-        #[inline]
-        fn next(&mut self) -> Option<Self::Item> {
-            unsafe {
-                let mut key: *mut ffi::PyObject = std::ptr::null_mut();
-                let mut hash: ffi::Py_hash_t = 0;
-                if ffi::_PySet_NextEntry(self.set.as_ptr(), &mut self.pos, &mut key, &mut hash) != 0
-                {
-                    // _PySet_NextEntry returns borrowed object; for safety must make owned (see #890)
-                    Some(key.assume_borrowed(self.set.py()).to_owned())
-                } else {
-                    None
-                }
-            }
-        }
+impl<'py> Iterator for BoundFrozenSetIterator<'py> {
+    type Item = Bound<'py, super::PyAny>;
 
-        #[inline]
-        fn size_hint(&self) -> (usize, Option<usize>) {
-            let len = self.len();
-            (len, Some(len))
-        }
+    /// Advances the iterator and returns the next value.
+    fn next(&mut self) -> Option<Self::Item> {
+        self.remaining = self.remaining.saturating_sub(1);
+        self.it.next().map(Result::unwrap)
     }
 
-    impl<'py> ExactSizeIterator for BoundFrozenSetIterator<'py> {
-        fn len(&self) -> usize {
-            self.set.len().saturating_sub(self.pos as usize)
-        }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.remaining, Some(self.remaining))
     }
+}
 
-    impl<'py> ExactSizeIterator for PyFrozenSetIterator<'py> {
-        fn len(&self) -> usize {
-            self.0.len()
-        }
+impl<'py> ExactSizeIterator for BoundFrozenSetIterator<'py> {
+    fn len(&self) -> usize {
+        self.remaining
     }
 }
 
-pub use impl_::*;
-
 #[inline]
 pub(crate) fn new_from_iter<T: ToPyObject>(
     py: Python<'_>,
@@ -387,7 +345,6 @@ mod tests {
     }
 
     #[test]
-    #[cfg(not(Py_LIMITED_API))]
     fn test_frozenset_iter_size_hint() {
         Python::with_gil(|py| {
             let set = PyFrozenSet::new(py, &[1]).unwrap();
@@ -402,18 +359,6 @@ mod tests {
         });
     }
 
-    #[test]
-    #[cfg(Py_LIMITED_API)]
-    fn test_frozenset_iter_size_hint() {
-        Python::with_gil(|py| {
-            let set = PyFrozenSet::new(py, &[1]).unwrap();
-            let iter = set.iter();
-
-            // No known bounds
-            assert_eq!(iter.size_hint(), (0, None));
-        });
-    }
-
     #[test]
     fn test_frozenset_builder() {
         use super::PyFrozenSetBuilder;

diff --git a/src/types/set.rs b/src/types/set.rs
@@ -1,4 +1,3 @@
-#[cfg(Py_LIMITED_API)]
 use crate::types::PyIterator;
 use crate::{
     err::{self, PyErr, PyResult},
@@ -277,6 +276,12 @@ impl<'py> Iterator for PySetIterator<'py> {
     }
 }
 
+impl ExactSizeIterator for PySetIterator<'_> {
+    fn len(&self) -> usize {
+        self.0.len()
+    }
+}
+
 impl<'py> IntoIterator for &'py PySet {
     type Item = &'py PyAny;
     type IntoIter = PySetIterator<'py>;
@@ -304,104 +309,43 @@ impl<'py> IntoIterator for Bound<'py, PySet> {
     }
 }
 
-#[cfg(Py_LIMITED_API)]
-mod impl_ {
-    use super::*;
-
-    /// PyO3 implementation of an iterator for a Python `set` object.
-    pub struct BoundSetIterator<'p> {
-        it: Bound<'p, PyIterator>,
-    }
-
-    impl<'py> BoundSetIterator<'py> {
-        pub(super) fn new(set: Bound<'py, PySet>) -> Self {
-            Self {
-                it: PyIterator::from_bound_object(&set).unwrap(),
-            }
-        }
-    }
-
-    impl<'py> Iterator for BoundSetIterator<'py> {
-        type Item = Bound<'py, super::PyAny>;
-
-        /// Advances the iterator and returns the next value.
-        ///
-        /// # Panics
-        ///
-        /// If PyO3 detects that the set is mutated during iteration, it will panic.
-        #[inline]
-        fn next(&mut self) -> Option<Self::Item> {
-            self.it.next().map(Result::unwrap)
-        }
-    }
+/// PyO3 implementation of an iterator for a Python `set` object.
+pub struct BoundSetIterator<'p> {
+    it: Bound<'p, PyIterator>,
+    // Remaining elements in the set. This is fine to store because
+    // Python will error if the set changes size during iteration.
+    remaining: usize,
 }
 
-#[cfg(not(Py_LIMITED_API))]
-mod impl_ {
-    use super::*;
-
-    /// PyO3 implementation of an iterator for a Python `set` object.
-    pub struct BoundSetIterator<'py> {
-        set: Bound<'py, super::PySet>,
-        pos: ffi::Py_ssize_t,
-        used: ffi::Py_ssize_t,
-    }
-
-    impl<'py> BoundSetIterator<'py> {
-        pub(super) fn new(set: Bound<'py, PySet>) -> Self {
-            let used = unsafe { ffi::PySet_Size(set.as_ptr()) };
-            BoundSetIterator { set, pos: 0, used }
+impl<'py> BoundSetIterator<'py> {
+    pub(super) fn new(set: Bound<'py, PySet>) -> Self {
+        Self {
+            it: PyIterator::from_bound_object(&set).unwrap(),
+            remaining: set.len(),
         }
     }
+}
 
-    impl<'py> Iterator for BoundSetIterator<'py> {
-        type Item = Bound<'py, super::PyAny>;
-
-        /// Advances the iterator and returns the next value.
-        ///
-        /// # Panics
-        ///
-        /// If PyO3 detects that the set is mutated during iteration, it will panic.
-        #[inline]
-        fn next(&mut self) -> Option<Self::Item> {
-            unsafe {
-                let len = ffi::PySet_Size(self.set.as_ptr());
-                assert_eq!(self.used, len, "Set changed size during iteration");
-
-                let mut key: *mut ffi::PyObject = std::ptr::null_mut();
-                let mut hash: ffi::Py_hash_t = 0;
-                if ffi::_PySet_NextEntry(self.set.as_ptr(), &mut self.pos, &mut key, &mut hash) != 0
-                {
-                    // _PySet_NextEntry returns borrowed object; for safety must make owned (see #890)
-                    Some(key.assume_borrowed(self.set.py()).to_owned())
-                } else {
-                    None
-                }
-            }
-        }
+impl<'py> Iterator for BoundSetIterator<'py> {
+    type Item = Bound<'py, super::PyAny>;
 
-        #[inline]
-        fn size_hint(&self) -> (usize, Option<usize>) {
-            let len = self.len();
-            (len, Some(len))
-        }
+    /// Advances the iterator and returns the next value.
+    fn next(&mut self) -> Option<Self::Item> {
+        self.remaining = self.remaining.saturating_sub(1);
+        self.it.next().map(Result::unwrap)
     }
 
-    impl<'py> ExactSizeIterator for BoundSetIterator<'py> {
-        fn len(&self) -> usize {
-            self.set.len().saturating_sub(self.pos as usize)
-        }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.remaining, Some(self.remaining))
     }
+}
 
-    impl<'py> ExactSizeIterator for PySetIterator<'py> {
-        fn len(&self) -> usize {
-            self.0.len()
-        }
+impl<'py> ExactSizeIterator for BoundSetIterator<'py> {
+    fn len(&self) -> usize {
+        self.remaining
     }
 }
 
-pub use impl_::*;
-
 #[inline]
 pub(crate) fn new_from_iter<T: ToPyObject>(
     py: Python<'_>,
@@ -571,7 +515,6 @@ mod tests {
     }
 
     #[test]
-    #[cfg(not(Py_LIMITED_API))]
     fn test_set_iter_size_hint() {
         Python::with_gil(|py| {
             let set = PySet::new(py, &[1]).unwrap();
@@ -585,16 +528,4 @@ mod tests {
             assert_eq!(iter.size_hint(), (0, Some(0)));
         });
     }
-
-    #[test]
-    #[cfg(Py_LIMITED_API)]
-    fn test_set_iter_size_hint() {
-        Python::with_gil(|py| {
-            let set = PySet::new(py, &[1]).unwrap();
-            let iter = set.iter();
-
-            // No known bounds
-            assert_eq!(iter.size_hint(), (0, None));
-        });
-    }
 }