diff --git a/CHANGELOG.md b/CHANGELOG.md index bd8c315ba0a..e0b3814e6ba 100644 --- a/CHANGELOG.md +++ b/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 for T` where `U: FromPy` 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) diff --git a/src/ffi/setobject.rs b/src/ffi/setobject.rs index d1e4dbd703e..605e1a9d68d 100644 --- 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 index c62ef734e32..bb75b9e557f 100644 --- a/src/types/set.rs +++ b/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 { + 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 ToPyObject for collections::HashSet 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::().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::().unwrap()); } + + // intoiterator iteration + for el in set { + assert_eq!(1i32, el.extract::().unwrap()); + } } }