pyo3/types/
frozenset.rs

1use crate::types::PyIterator;
2use crate::{
3    err::{self, PyErr, PyResult},
4    ffi,
5    ffi_ptr_ext::FfiPtrExt,
6    py_result_ext::PyResultExt,
7    types::any::PyAnyMethods,
8    Bound, PyAny, Python,
9};
10use crate::{Borrowed, BoundObject, IntoPyObject, IntoPyObjectExt};
11use std::ptr;
12
13/// Allows building a Python `frozenset` one item at a time
14pub struct PyFrozenSetBuilder<'py> {
15    py_frozen_set: Bound<'py, PyFrozenSet>,
16}
17
18impl<'py> PyFrozenSetBuilder<'py> {
19    /// Create a new `FrozenSetBuilder`.
20    /// Since this allocates a `PyFrozenSet` internally it may
21    /// panic when running out of memory.
22    pub fn new(py: Python<'py>) -> PyResult<PyFrozenSetBuilder<'py>> {
23        Ok(PyFrozenSetBuilder {
24            py_frozen_set: PyFrozenSet::empty(py)?,
25        })
26    }
27
28    /// Adds an element to the set.
29    pub fn add<K>(&mut self, key: K) -> PyResult<()>
30    where
31        K: IntoPyObject<'py>,
32    {
33        fn inner(frozenset: &Bound<'_, PyFrozenSet>, key: Borrowed<'_, '_, PyAny>) -> PyResult<()> {
34            err::error_on_minusone(frozenset.py(), unsafe {
35                ffi::PySet_Add(frozenset.as_ptr(), key.as_ptr())
36            })
37        }
38
39        inner(
40            &self.py_frozen_set,
41            key.into_pyobject(self.py_frozen_set.py())
42                .map_err(Into::into)?
43                .into_any()
44                .as_borrowed(),
45        )
46    }
47
48    /// Finish building the set and take ownership of its current value
49    pub fn finalize(self) -> Bound<'py, PyFrozenSet> {
50        self.py_frozen_set
51    }
52}
53
54/// Represents a  Python `frozenset`.
55///
56/// Values of this type are accessed via PyO3's smart pointers, e.g. as
57/// [`Py<PyFrozenSet>`][crate::Py] or [`Bound<'py, PyFrozenSet>`][Bound].
58///
59/// For APIs available on `frozenset` objects, see the [`PyFrozenSetMethods`] trait which is implemented for
60/// [`Bound<'py, PyFrozenSet>`][Bound].
61#[repr(transparent)]
62pub struct PyFrozenSet(PyAny);
63
64#[cfg(not(any(PyPy, GraalPy)))]
65pyobject_subclassable_native_type!(PyFrozenSet, crate::ffi::PySetObject);
66#[cfg(not(any(PyPy, GraalPy)))]
67pyobject_native_type!(
68    PyFrozenSet,
69    ffi::PySetObject,
70    pyobject_native_static_type_object!(ffi::PyFrozenSet_Type),
71    #checkfunction=ffi::PyFrozenSet_Check
72);
73
74#[cfg(any(PyPy, GraalPy))]
75pyobject_native_type_core!(
76    PyFrozenSet,
77    pyobject_native_static_type_object!(ffi::PyFrozenSet_Type),
78    #checkfunction=ffi::PyFrozenSet_Check
79);
80
81impl PyFrozenSet {
82    /// Creates a new frozenset.
83    ///
84    /// May panic when running out of memory.
85    #[inline]
86    pub fn new<'py, T>(
87        py: Python<'py>,
88        elements: impl IntoIterator<Item = T>,
89    ) -> PyResult<Bound<'py, PyFrozenSet>>
90    where
91        T: IntoPyObject<'py>,
92    {
93        try_new_from_iter(py, elements)
94    }
95
96    /// Creates a new empty frozen set
97    pub fn empty(py: Python<'_>) -> PyResult<Bound<'_, PyFrozenSet>> {
98        unsafe {
99            ffi::PyFrozenSet_New(ptr::null_mut())
100                .assume_owned_or_err(py)
101                .downcast_into_unchecked()
102        }
103    }
104}
105
106/// Implementation of functionality for [`PyFrozenSet`].
107///
108/// These methods are defined for the `Bound<'py, PyFrozenSet>` smart pointer, so to use method call
109/// syntax these methods are separated into a trait, because stable Rust does not yet support
110/// `arbitrary_self_types`.
111#[doc(alias = "PyFrozenSet")]
112pub trait PyFrozenSetMethods<'py>: crate::sealed::Sealed {
113    /// Returns the number of items in the set.
114    ///
115    /// This is equivalent to the Python expression `len(self)`.
116    fn len(&self) -> usize;
117
118    /// Checks if set is empty.
119    fn is_empty(&self) -> bool {
120        self.len() == 0
121    }
122
123    /// Determines if the set contains the specified key.
124    ///
125    /// This is equivalent to the Python expression `key in self`.
126    fn contains<K>(&self, key: K) -> PyResult<bool>
127    where
128        K: IntoPyObject<'py>;
129
130    /// Returns an iterator of values in this set.
131    fn iter(&self) -> BoundFrozenSetIterator<'py>;
132}
133
134impl<'py> PyFrozenSetMethods<'py> for Bound<'py, PyFrozenSet> {
135    #[inline]
136    fn len(&self) -> usize {
137        unsafe { ffi::PySet_Size(self.as_ptr()) as usize }
138    }
139
140    fn contains<K>(&self, key: K) -> PyResult<bool>
141    where
142        K: IntoPyObject<'py>,
143    {
144        fn inner(
145            frozenset: &Bound<'_, PyFrozenSet>,
146            key: Borrowed<'_, '_, PyAny>,
147        ) -> PyResult<bool> {
148            match unsafe { ffi::PySet_Contains(frozenset.as_ptr(), key.as_ptr()) } {
149                1 => Ok(true),
150                0 => Ok(false),
151                _ => Err(PyErr::fetch(frozenset.py())),
152            }
153        }
154
155        let py = self.py();
156        inner(
157            self,
158            key.into_pyobject_or_pyerr(py)?.into_any().as_borrowed(),
159        )
160    }
161
162    fn iter(&self) -> BoundFrozenSetIterator<'py> {
163        BoundFrozenSetIterator::new(self.clone())
164    }
165}
166
167impl<'py> IntoIterator for Bound<'py, PyFrozenSet> {
168    type Item = Bound<'py, PyAny>;
169    type IntoIter = BoundFrozenSetIterator<'py>;
170
171    /// Returns an iterator of values in this set.
172    fn into_iter(self) -> Self::IntoIter {
173        BoundFrozenSetIterator::new(self)
174    }
175}
176
177impl<'py> IntoIterator for &Bound<'py, PyFrozenSet> {
178    type Item = Bound<'py, PyAny>;
179    type IntoIter = BoundFrozenSetIterator<'py>;
180
181    /// Returns an iterator of values in this set.
182    fn into_iter(self) -> Self::IntoIter {
183        self.iter()
184    }
185}
186
187/// PyO3 implementation of an iterator for a Python `frozenset` object.
188pub struct BoundFrozenSetIterator<'p> {
189    it: Bound<'p, PyIterator>,
190    // Remaining elements in the frozenset
191    remaining: usize,
192}
193
194impl<'py> BoundFrozenSetIterator<'py> {
195    pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
196        Self {
197            it: PyIterator::from_object(&set).unwrap(),
198            remaining: set.len(),
199        }
200    }
201}
202
203impl<'py> Iterator for BoundFrozenSetIterator<'py> {
204    type Item = Bound<'py, super::PyAny>;
205
206    /// Advances the iterator and returns the next value.
207    fn next(&mut self) -> Option<Self::Item> {
208        self.remaining = self.remaining.saturating_sub(1);
209        self.it.next().map(Result::unwrap)
210    }
211
212    fn size_hint(&self) -> (usize, Option<usize>) {
213        (self.remaining, Some(self.remaining))
214    }
215
216    #[inline]
217    fn count(self) -> usize
218    where
219        Self: Sized,
220    {
221        self.len()
222    }
223}
224
225impl ExactSizeIterator for BoundFrozenSetIterator<'_> {
226    fn len(&self) -> usize {
227        self.remaining
228    }
229}
230
231#[inline]
232pub(crate) fn try_new_from_iter<'py, T>(
233    py: Python<'py>,
234    elements: impl IntoIterator<Item = T>,
235) -> PyResult<Bound<'py, PyFrozenSet>>
236where
237    T: IntoPyObject<'py>,
238{
239    let set = unsafe {
240        // We create the  `Py` pointer because its Drop cleans up the set if user code panics.
241        ffi::PyFrozenSet_New(std::ptr::null_mut())
242            .assume_owned_or_err(py)?
243            .downcast_into_unchecked()
244    };
245    let ptr = set.as_ptr();
246
247    for e in elements {
248        let obj = e.into_pyobject_or_pyerr(py)?;
249        err::error_on_minusone(py, unsafe { ffi::PySet_Add(ptr, obj.as_ptr()) })?;
250    }
251
252    Ok(set)
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    #[test]
260    fn test_frozenset_new_and_len() {
261        Python::with_gil(|py| {
262            let set = PyFrozenSet::new(py, [1]).unwrap();
263            assert_eq!(1, set.len());
264
265            let v = vec![1];
266            assert!(PyFrozenSet::new(py, &[v]).is_err());
267        });
268    }
269
270    #[test]
271    fn test_frozenset_empty() {
272        Python::with_gil(|py| {
273            let set = PyFrozenSet::empty(py).unwrap();
274            assert_eq!(0, set.len());
275            assert!(set.is_empty());
276        });
277    }
278
279    #[test]
280    fn test_frozenset_contains() {
281        Python::with_gil(|py| {
282            let set = PyFrozenSet::new(py, [1]).unwrap();
283            assert!(set.contains(1).unwrap());
284        });
285    }
286
287    #[test]
288    fn test_frozenset_iter() {
289        Python::with_gil(|py| {
290            let set = PyFrozenSet::new(py, [1]).unwrap();
291
292            for el in set {
293                assert_eq!(1i32, el.extract::<i32>().unwrap());
294            }
295        });
296    }
297
298    #[test]
299    fn test_frozenset_iter_bound() {
300        Python::with_gil(|py| {
301            let set = PyFrozenSet::new(py, [1]).unwrap();
302
303            for el in &set {
304                assert_eq!(1i32, el.extract::<i32>().unwrap());
305            }
306        });
307    }
308
309    #[test]
310    fn test_frozenset_iter_size_hint() {
311        Python::with_gil(|py| {
312            let set = PyFrozenSet::new(py, [1]).unwrap();
313            let mut iter = set.iter();
314
315            // Exact size
316            assert_eq!(iter.len(), 1);
317            assert_eq!(iter.size_hint(), (1, Some(1)));
318            iter.next();
319            assert_eq!(iter.len(), 0);
320            assert_eq!(iter.size_hint(), (0, Some(0)));
321        });
322    }
323
324    #[test]
325    fn test_frozenset_builder() {
326        use super::PyFrozenSetBuilder;
327
328        Python::with_gil(|py| {
329            let mut builder = PyFrozenSetBuilder::new(py).unwrap();
330
331            // add an item
332            builder.add(1).unwrap();
333            builder.add(2).unwrap();
334            builder.add(2).unwrap();
335
336            // finalize it
337            let set = builder.finalize();
338
339            assert!(set.contains(1).unwrap());
340            assert!(set.contains(2).unwrap());
341            assert!(!set.contains(3).unwrap());
342        });
343    }
344
345    #[test]
346    fn test_iter_count() {
347        Python::with_gil(|py| {
348            let set = PyFrozenSet::new(py, vec![1, 2, 3]).unwrap();
349            assert_eq!(set.iter().count(), 3);
350        })
351    }
352}
pyo3/types/frozenset.rs

pyo3/types/
frozenset.rs
