Use _PySet_NextEntry

PyO3 · Jan 8, 2020 · f7a4fba · f7a4fba
1 parent ec79285
commit f7a4fba
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 11 deletions.
diff --git a/src/ffi/setobject.rs b/src/ffi/setobject.rs
@@ -60,4 +60,13 @@ extern "C" {
     pub fn PySet_Add(set: *mut PyObject, key: *mut PyObject) -> c_int;
     #[cfg_attr(PyPy, link_name = "PyPySet_Pop")]
     pub fn PySet_Pop(set: *mut PyObject) -> *mut PyObject;
+
+    #[cfg(not(Py_LIMITED_API))]
+    #[cfg_attr(PyPy, link_name = "_PySet_NextEntry")]
+    pub fn _PySet_NextEntry(
+        set: *mut PyObject,
+        pos: *mut Py_ssize_t,
+        key: *mut *mut PyObject,
+        hash: *mut super::Py_hash_t,
+    ) -> c_int;
 }
diff --git a/src/types/set.rs b/src/types/set.rs
@@ -6,8 +6,7 @@ use crate::ffi;
 use crate::instance::PyNativeType;
 use crate::internal_tricks::Unsendable;
 use crate::object::PyObject;
-use crate::objectprotocol::ObjectProtocol;
-use crate::types::{PyAny, PyIterator};
+use crate::types::PyAny;
 use crate::AsPyPointer;
 use crate::Python;
 use crate::{ToBorrowedObject, ToPyObject};
@@ -98,12 +97,40 @@ impl PySet {
     }
 }
 
+#[cfg(not(Py_LIMITED_API))]
+pub struct PySetIterator<'py> {
+    set: &'py super::PyAny,
+    pos: isize,
+}
+
+#[cfg(not(Py_LIMITED_API))]
+impl<'py> Iterator for PySetIterator<'py> {
+    type Item = &'py super::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 {
+                Some(self.set.py().from_borrowed_ptr(key))
+            } else {
+                None
+            }
+        }
+    }
+}
+
+#[cfg(not(Py_LIMITED_API))]
 impl<'a> std::iter::IntoIterator for &'a PySet {
-    type Item = PyResult<&'a PyAny>;
-    type IntoIter = PyIterator<'a>;
+    type Item = &'a PyAny;
+    type IntoIter = PySetIterator<'a>;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.iter().expect("Failed to get set iterator")
+        PySetIterator {
+            set: self.as_ref(),
+            pos: 0,
+        }
     }
 }
 
@@ -178,13 +205,16 @@ impl PyFrozenSet {
         })
     }
 }
-
+#[cfg(not(Py_LIMITED_API))]
 impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
-    type Item = PyResult<&'a PyAny>;
-    type IntoIter = PyIterator<'a>;
+    type Item = &'a PyAny;
+    type IntoIter = PySetIterator<'a>;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.iter().expect("Failed to get frozen set iterator")
+        PySetIterator {
+            set: self.as_ref(),
+            pos: 0,
+        }
     }
 }
 
@@ -294,7 +324,7 @@ mod test {
 
         // intoiterator iteration
         for el in set {
-            assert_eq!(1i32, el.unwrap().extract().unwrap());
+            assert_eq!(1i32, el.extract().unwrap());
         }
     }
 
@@ -340,7 +370,7 @@ mod test {
 
         // intoiterator iteration
         for el in set {
-            assert_eq!(1i32, el.unwrap().extract::<i32>().unwrap());
+            assert_eq!(1i32, el.extract::<i32>().unwrap());
         }
     }
 }