Merge pull request #716 from davidhewitt/set-iterator

Implement IntoIterator for PySet and PyFrozenSet

CHANGELOG.md

@@ -12,6 +12,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 * The blanket implementations for `FromPyObject` for `&T` and `&mut T` are no longer specializable. Implement `PyTryFrom` for your type to control the behavior of `FromPyObject::extract()` for your types.
 * The implementation for `IntoPy<U> for T` where `U: FromPy<T>` is no longer specializable. Control the behavior of this via the implementation of `FromPy`.
 
+### Added
+
+* Implemented `IntoIterator` for `PySet` and `PyFrozenSet`. [#716](https://github.com/PyO3/pyo3/pull/716)
+
 ## [0.8.5]
 
 * Support for `#[name = "foo"]` attribute for `#[pyfunction]` and in `#[pymethods]`. [#692](https://github.com/PyO3/pyo3/pull/692)

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

src/types/set.rs

@@ -6,6 +6,7 @@ use crate::ffi;
 use crate::instance::PyNativeType;
 use crate::internal_tricks::Unsendable;
 use crate::object::PyObject;
+use crate::types::PyAny;
 use crate::AsPyPointer;
 use crate::Python;
 use crate::{ToBorrowedObject, ToPyObject};
@@ -96,6 +97,43 @@ 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 = &'a PyAny;
+    type IntoIter = PySetIterator<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        PySetIterator {
+            set: self.as_ref(),
+            pos: 0,
+        }
+    }
+}
+
 impl<T> ToPyObject for collections::HashSet<T>
 where
     T: hash::Hash + Eq + ToPyObject,
@@ -167,6 +205,18 @@ impl PyFrozenSet {
         })
     }
 }
+#[cfg(not(Py_LIMITED_API))]
+impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
+    type Item = &'a PyAny;
+    type IntoIter = PySetIterator<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        PySetIterator {
+            set: self.as_ref(),
+            pos: 0,
+        }
+    }
+}
 
 #[cfg(test)]
 mod test {
@@ -267,9 +317,15 @@ mod test {
         let py = gil.python();
 
         let set = PySet::new(py, &[1]).unwrap();
+        // objectprotocol iteration
         for el in set.iter().unwrap() {
             assert_eq!(1i32, el.unwrap().extract::<i32>().unwrap());
         }
+
+        // intoiterator iteration
+        for el in set {
+            assert_eq!(1i32, el.extract().unwrap());
+        }
     }
 
     #[test]
@@ -306,8 +362,15 @@ mod test {
         let py = gil.python();
 
         let set = PyFrozenSet::new(py, &[1]).unwrap();
+
+        // objectprotocol iteration
         for el in set.iter().unwrap() {
             assert_eq!(1i32, el.unwrap().extract::<i32>().unwrap());
         }
+
+        // intoiterator iteration
+        for el in set {
+            assert_eq!(1i32, el.extract::<i32>().unwrap());
+        }
     }
 }