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

pyo3/types/
frozenset.rs
