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

pyo3/types/
frozenset.rs
