luxolus
/
pyo3
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

always use a Python iterator for sets and frozensets (#3849) 

* always use a Python iterator for sets and frozensets

* add newsfragment
This commit is contained in:
David Hewitt 2024-02-17 10:57:53 +00:00 committed by GitHub
parent 65cf5808d9
commit eb90b81d44
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 64 additions and 186 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
newsfragments/3849.added.md Normal file
View File

@ -0,0 +1 @@
Implement `ExactSizeIterator` for `set` and `frozenset` iterators on `abi3` feature.

1
newsfragments/3849.changed.md Normal file
View File

@ -0,0 +1 @@
`PySet` and `PyFrozenSet` iterators now always iterate the equivalent of `iter(set)`. (A "fast path" with no noticeable performance benefit was removed.)

117
src/types/frozenset.rs
View File

@ -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,89 +230,41 @@ impl<'py> IntoIterator for Bound<'py, PyFrozenSet> {
}
}
#[cfg(Py_LIMITED_API)]
mod impl_ {
use super::*;
/// 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,
}
/// 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)
impl<'py> BoundFrozenSetIterator<'py> {
pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
Self {
it: PyIterator::from_bound_object(&set).unwrap(),
remaining: set.len(),
}
}
}
#[cfg(not(Py_LIMITED_API))]
mod impl_ {
use super::*;
impl<'py> Iterator for BoundFrozenSetIterator<'py> {
type Item = Bound<'py, super::PyAny>;
/// PyO3 implementation of an iterator for a Python `frozenset` object.
pub struct BoundFrozenSetIterator<'py> {
set: Bound<'py, PyFrozenSet>,
pos: ffi::Py_ssize_t,
/// 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> BoundFrozenSetIterator<'py> {
pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
Self { set, pos: 0 }
}
}
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
}
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.len();
(len, Some(len))
}
}
impl<'py> ExactSizeIterator for BoundFrozenSetIterator<'py> {
fn len(&self) -> usize {
self.set.len().saturating_sub(self.pos as usize)
}
}
impl<'py> ExactSizeIterator for PyFrozenSetIterator<'py> {
fn len(&self) -> usize {
self.0.len()
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.remaining, Some(self.remaining))
}
}
pub use impl_::*;
impl<'py> ExactSizeIterator for BoundFrozenSetIterator<'py> {
fn len(&self) -> usize {
self.remaining
}
}
#[inline]
pub(crate) fn new_from_iter<T: ToPyObject>(
@ -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;

131
src/types/set.rs
View File

@ -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,103 +309,42 @@ 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>,
// Remaining elements in the set. This is fine to store because
// Python will error if the set changes size during iteration.
remaining: usize,
}
/// 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)
impl<'py> BoundSetIterator<'py> {
pub(super) fn new(set: Bound<'py, PySet>) -> Self {
Self {
it: PyIterator::from_bound_object(&set).unwrap(),
remaining: set.len(),
}
}
}
#[cfg(not(Py_LIMITED_API))]
mod impl_ {
use super::*;
impl<'py> Iterator for BoundSetIterator<'py> {
type Item = Bound<'py, super::PyAny>;
/// 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,
/// 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> 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> 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
}
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.len();
(len, Some(len))
}
}
impl<'py> ExactSizeIterator for BoundSetIterator<'py> {
fn len(&self) -> usize {
self.set.len().saturating_sub(self.pos as usize)
}
}
impl<'py> ExactSizeIterator for PySetIterator<'py> {
fn len(&self) -> usize {
self.0.len()
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.remaining, Some(self.remaining))
}
}
pub use impl_::*;
impl<'py> ExactSizeIterator for BoundSetIterator<'py> {
fn len(&self) -> usize {
self.remaining
}
}
#[inline]
pub(crate) fn new_from_iter<T: ToPyObject>(
@ -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));
});
}
}