From 8a28a69c3d51972d9bdc24141fcfe11eec0ac27b Mon Sep 17 00:00:00 2001 From: David Hewitt Date: Mon, 27 Nov 2023 13:55:53 +0000 Subject: [PATCH] implement `PyFrozenSetMethods` --- src/prelude.rs | 1 + src/types/frozenset.rs | 199 +++++++++++++++++++++++++++++++++-------- src/types/mod.rs | 4 +- 3 files changed, 164 insertions(+), 40 deletions(-) diff --git a/src/prelude.rs b/src/prelude.rs index beaa385f..9be5dc0d 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -31,6 +31,7 @@ pub use crate::types::bytearray::PyByteArrayMethods; pub use crate::types::bytes::PyBytesMethods; pub use crate::types::dict::PyDictMethods; pub use crate::types::float::PyFloatMethods; +pub use crate::types::frozenset::PyFrozenSetMethods; pub use crate::types::list::PyListMethods; pub use crate::types::mapping::PyMappingMethods; pub use crate::types::sequence::PySequenceMethods; diff --git a/src/types/frozenset.rs b/src/types/frozenset.rs index bbcd7102..777abf6f 100644 --- a/src/types/frozenset.rs +++ b/src/types/frozenset.rs @@ -1,11 +1,11 @@ +#[cfg(not(Py_LIMITED_API))] +use crate::ffi_ptr_ext::FfiPtrExt; #[cfg(Py_LIMITED_API)] use crate::types::PyIterator; use crate::{ err::{self, PyErr, PyResult}, - Py, PyObject, + ffi, Bound, Py, PyAny, PyNativeType, PyObject, Python, ToPyObject, }; -use crate::{ffi, PyAny, Python, ToPyObject}; - use std::ptr; /// Allows building a Python `frozenset` one item at a time @@ -83,12 +83,12 @@ impl PyFrozenSet { /// This is equivalent to len(p) on a set. #[inline] pub fn len(&self) -> usize { - unsafe { ffi::PySet_Size(self.as_ptr()) as usize } + self.as_borrowed().len() } /// Check if set is empty. pub fn is_empty(&self) -> bool { - self.len() == 0 + self.as_borrowed().is_empty() } /// Determine if the set contains the specified key. @@ -97,7 +97,54 @@ impl PyFrozenSet { where K: ToPyObject, { - fn inner(frozenset: &PyFrozenSet, key: PyObject) -> PyResult { + self.as_borrowed().contains(key) + } + + /// Returns an iterator of values in this frozen set. + pub fn iter(&self) -> PyFrozenSetIterator<'_> { + PyFrozenSetIterator(BoundFrozenSetIterator::new(self.as_borrowed().to_owned())) + } +} + +/// Implementation of functionality for [`PyFrozenSet`]. +/// +/// These methods are defined for the `Bound<'py, PyFrozenSet>` smart pointer, so to use method call +/// syntax these methods are separated into a trait, because stable Rust does not yet support +/// `arbitrary_self_types`. +#[doc(alias = "PyFrozenSet")] +pub trait PyFrozenSetMethods<'py> { + /// Returns the number of items in the set. + /// + /// This is equivalent to the Python expression `len(self)`. + fn len(&self) -> usize; + + /// Checks if set is empty. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Determines if the set contains the specified key. + /// + /// This is equivalent to the Python expression `key in self`. + fn contains(&self, key: K) -> PyResult + where + K: ToPyObject; + + /// Returns an iterator of values in this set. + fn iter(&self) -> BoundFrozenSetIterator<'py>; +} + +impl<'py> PyFrozenSetMethods<'py> for Bound<'py, PyFrozenSet> { + #[inline] + fn len(&self) -> usize { + unsafe { ffi::PySet_Size(self.as_ptr()) as usize } + } + + fn contains(&self, key: K) -> PyResult + where + K: ToPyObject, + { + fn inner(frozenset: &Bound<'_, PyFrozenSet>, key: Bound<'_, PyAny>) -> PyResult { match unsafe { ffi::PySet_Contains(frozenset.as_ptr(), key.as_ptr()) } { 1 => Ok(true), 0 => Ok(false), @@ -105,12 +152,58 @@ impl PyFrozenSet { } } - inner(self, key.to_object(self.py())) + let py = self.py(); + inner(self, key.to_object(py).into_bound(py)) } - /// Returns an iterator of values in this frozen set. - pub fn iter(&self) -> PyFrozenSetIterator<'_> { - IntoIterator::into_iter(self) + fn iter(&self) -> BoundFrozenSetIterator<'py> { + BoundFrozenSetIterator::new(self.clone()) + } +} + +/// PyO3 implementation of an iterator for a Python `frozenset` object. +pub struct PyFrozenSetIterator<'py>(BoundFrozenSetIterator<'py>); + +impl<'py> Iterator for PyFrozenSetIterator<'py> { + type Item = &'py super::PyAny; + + /// Advances the iterator and returns the next value. + #[inline] + fn next(&mut self) -> Option { + self.0.next().map(Bound::into_gil_ref) + } + + fn size_hint(&self) -> (usize, Option) { + self.0.size_hint() + } +} + +impl<'py> IntoIterator for &'py PyFrozenSet { + type Item = &'py PyAny; + type IntoIter = PyFrozenSetIterator<'py>; + /// Returns an iterator of values in this set. + fn into_iter(self) -> Self::IntoIter { + PyFrozenSetIterator(BoundFrozenSetIterator::new(self.as_borrowed().to_owned())) + } +} + +impl<'py> IntoIterator for Bound<'py, PyFrozenSet> { + type Item = Bound<'py, PyAny>; + type IntoIter = BoundFrozenSetIterator<'py>; + + /// Returns an iterator of values in this set. + fn into_iter(self) -> Self::IntoIter { + BoundFrozenSetIterator::new(self) + } +} + +impl<'py> IntoIterator for &'_ Bound<'py, PyFrozenSet> { + type Item = Bound<'py, PyAny>; + type IntoIter = BoundFrozenSetIterator<'py>; + + /// Returns an iterator of values in this set. + fn into_iter(self) -> Self::IntoIter { + BoundFrozenSetIterator::new(self.clone()) } } @@ -118,25 +211,23 @@ impl PyFrozenSet { mod impl_ { use super::*; - impl<'a> std::iter::IntoIterator for &'a PyFrozenSet { - type Item = &'a PyAny; - type IntoIter = PyFrozenSetIterator<'a>; + /// PyO3 implementation of an iterator for a Python `set` object. + pub struct BoundFrozenSetIterator<'p> { + it: Bound<'p, PyIterator>, + } - fn into_iter(self) -> Self::IntoIter { - PyFrozenSetIterator { - it: PyIterator::from_object(self).unwrap(), + impl<'py> BoundFrozenSetIterator<'py> { + pub(super) fn new(frozenset: Bound<'py, PyFrozenSet>) -> Self { + Self { + it: PyIterator::from_object2(&frozenset).unwrap(), } } } - /// PyO3 implementation of an iterator for a Python `frozenset` object. - pub struct PyFrozenSetIterator<'p> { - it: &'p PyIterator, - } - - impl<'py> Iterator for PyFrozenSetIterator<'py> { - type Item = &'py super::PyAny; + 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.it.next().map(Result::unwrap) @@ -148,23 +239,20 @@ mod impl_ { mod impl_ { use super::*; - impl<'a> std::iter::IntoIterator for &'a PyFrozenSet { - type Item = &'a PyAny; - type IntoIter = PyFrozenSetIterator<'a>; - - fn into_iter(self) -> Self::IntoIter { - PyFrozenSetIterator { set: self, pos: 0 } - } - } - /// PyO3 implementation of an iterator for a Python `frozenset` object. - pub struct PyFrozenSetIterator<'py> { - set: &'py PyFrozenSet, + pub struct BoundFrozenSetIterator<'py> { + set: Bound<'py, PyFrozenSet>, pos: ffi::Py_ssize_t, } - impl<'py> Iterator for PyFrozenSetIterator<'py> { - type Item = &'py PyAny; + 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 { @@ -174,7 +262,7 @@ mod impl_ { 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(self.set.py().from_owned_ptr(ffi::_Py_NewRef(key))) + Some(key.assume_borrowed(self.set.py()).to_owned()) } else { None } @@ -188,11 +276,17 @@ mod impl_ { } } - impl<'py> ExactSizeIterator for PyFrozenSetIterator<'py> { + 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() + } + } } pub use impl_::*; @@ -243,6 +337,7 @@ mod tests { Python::with_gil(|py| { let set = PyFrozenSet::empty(py).unwrap(); assert_eq!(0, set.len()); + assert!(set.is_empty()); }); } @@ -271,6 +366,34 @@ 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(); + let mut iter = set.iter(); + + // Exact size + assert_eq!(iter.len(), 1); + assert_eq!(iter.size_hint(), (1, Some(1))); + iter.next(); + assert_eq!(iter.len(), 0); + assert_eq!(iter.size_hint(), (0, Some(0))); + }); + } + + #[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/mod.rs b/src/types/mod.rs index a8770d14..705d3666 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -77,7 +77,7 @@ pub use self::typeobject::PyType; /// In these cases the iterators are implemented by forwarding to [`PyIterator`]. pub mod iter { pub use super::dict::{BoundDictIterator, PyDictIterator}; - pub use super::frozenset::PyFrozenSetIterator; + pub use super::frozenset::{BoundFrozenSetIterator, PyFrozenSetIterator}; pub use super::list::{BoundListIterator, PyListIterator}; pub use super::set::{BoundSetIterator, PySetIterator}; pub use super::tuple::PyTupleIterator; @@ -288,7 +288,7 @@ mod ellipsis; pub(crate) mod float; #[cfg(all(not(Py_LIMITED_API), not(PyPy)))] mod frame; -mod frozenset; +pub(crate) mod frozenset; mod function; mod iterator; pub(crate) mod list;