pyo3/types/
frozenset.rs

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

pyo3/types/
frozenset.rs
