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};
9#[cfg(RustPython)]
10use crate::{
11    sync::PyOnceLock,
12    types::{PyType, PyTypeMethods},
13    Py,
14};
15use crate::{Borrowed, BoundObject, IntoPyObject, IntoPyObjectExt};
16use core::ptr;
17
18/// Allows building a Python `frozenset` one item at a time
19pub struct PyFrozenSetBuilder<'py> {
20    py_frozen_set: Bound<'py, PyFrozenSet>,
21}
22
23impl<'py> PyFrozenSetBuilder<'py> {
24    /// Create a new `FrozenSetBuilder`.
25    /// Since this allocates a `PyFrozenSet` internally it may
26    /// panic when running out of memory.
27    pub fn new(py: Python<'py>) -> PyResult<PyFrozenSetBuilder<'py>> {
28        Ok(PyFrozenSetBuilder {
29            py_frozen_set: PyFrozenSet::empty(py)?,
30        })
31    }
32
33    /// Adds an element to the set.
34    pub fn add<K>(&mut self, key: K) -> PyResult<()>
35    where
36        K: IntoPyObject<'py>,
37    {
38        fn inner(frozenset: &Bound<'_, PyFrozenSet>, key: Borrowed<'_, '_, PyAny>) -> PyResult<()> {
39            err::error_on_minusone(frozenset.py(), unsafe {
40                ffi::PySet_Add(frozenset.as_ptr(), key.as_ptr())
41            })
42        }
43
44        inner(
45            &self.py_frozen_set,
46            key.into_pyobject_or_pyerr(self.py_frozen_set.py())?
47                .into_any()
48                .as_borrowed(),
49        )
50    }
51
52    /// Finish building the set and take ownership of its current value
53    pub fn finalize(self) -> Bound<'py, PyFrozenSet> {
54        self.py_frozen_set
55    }
56}
57
58/// Represents a  Python `frozenset`.
59///
60/// Values of this type are accessed via PyO3's smart pointers, e.g. as
61/// [`Py<PyFrozenSet>`][crate::Py] or [`Bound<'py, PyFrozenSet>`][Bound].
62///
63/// For APIs available on `frozenset` objects, see the [`PyFrozenSetMethods`] trait which is implemented for
64/// [`Bound<'py, PyFrozenSet>`][Bound].
65#[repr(transparent)]
66pub struct PyFrozenSet(PyAny);
67
68#[cfg(not(any(PyPy, GraalPy)))]
69pyobject_subclassable_native_type!(PyFrozenSet, crate::ffi::PySetObject);
70#[cfg(all(not(any(PyPy, GraalPy)), not(RustPython)))]
71pyobject_native_type!(
72    PyFrozenSet,
73    ffi::PySetObject,
74    pyobject_native_static_type_object!(ffi::PyFrozenSet_Type),
75    "builtins",
76    "frozenset",
77    #checkfunction=ffi::PyFrozenSet_Check
78);
79
80#[cfg(all(not(any(PyPy, GraalPy)), RustPython))]
81pyobject_native_type!(
82    PyFrozenSet,
83    ffi::PySetObject,
84    |py| {
85        static TYPE: PyOnceLock<Py<PyType>> = PyOnceLock::new();
86        TYPE.import(py, "builtins", "frozenset").unwrap().as_type_ptr()
87    },
88    "builtins",
89    "frozenset",
90    #checkfunction=ffi::PyFrozenSet_Check
91);
92
93#[cfg(any(PyPy, GraalPy))]
94pyobject_native_type_core!(
95    PyFrozenSet,
96    pyobject_native_static_type_object!(ffi::PyFrozenSet_Type),
97    "builtins",
98    "frozenset",
99    #checkfunction=ffi::PyFrozenSet_Check
100);
101
102impl PyFrozenSet {
103    /// Creates a new frozenset.
104    ///
105    /// May panic when running out of memory.
106    #[inline]
107    pub fn new<'py, T>(
108        py: Python<'py>,
109        elements: impl IntoIterator<Item = T>,
110    ) -> PyResult<Bound<'py, PyFrozenSet>>
111    where
112        T: IntoPyObject<'py>,
113    {
114        let mut builder = PyFrozenSetBuilder::new(py)?;
115        for e in elements {
116            builder.add(e)?;
117        }
118        Ok(builder.finalize())
119    }
120
121    /// Creates a new empty frozen set
122    pub fn empty(py: Python<'_>) -> PyResult<Bound<'_, PyFrozenSet>> {
123        unsafe {
124            ffi::PyFrozenSet_New(ptr::null_mut())
125                .assume_owned_or_err(py)
126                .cast_into_unchecked()
127        }
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<'py>(Bound<'py, PyIterator>);
214
215impl<'py> BoundFrozenSetIterator<'py> {
216    pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
217        Self(PyIterator::from_object(&set).expect("frozenset should always be iterable"))
218    }
219}
220
221impl<'py> Iterator for BoundFrozenSetIterator<'py> {
222    type Item = Bound<'py, super::PyAny>;
223
224    /// Advances the iterator and returns the next value.
225    fn next(&mut self) -> Option<Self::Item> {
226        self.0
227            .next()
228            .map(|result| result.expect("frozenset iteration should be infallible"))
229    }
230
231    fn size_hint(&self) -> (usize, Option<usize>) {
232        let len = ExactSizeIterator::len(self);
233        (len, Some(len))
234    }
235
236    #[inline]
237    fn count(self) -> usize
238    where
239        Self: Sized,
240    {
241        self.len()
242    }
243}
244
245impl ExactSizeIterator for BoundFrozenSetIterator<'_> {
246    fn len(&self) -> usize {
247        self.0.size_hint().0
248    }
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254    use crate::types::PyAnyMethods as _;
255
256    #[test]
257    fn test_frozenset_new_and_len() {
258        Python::attach(|py| {
259            let set = PyFrozenSet::new(py, [1]).unwrap();
260            assert_eq!(1, set.len());
261
262            let v = vec![1];
263            assert!(PyFrozenSet::new(py, &[v]).is_err());
264        });
265    }
266
267    #[test]
268    fn test_frozenset_empty() {
269        Python::attach(|py| {
270            let set = PyFrozenSet::empty(py).unwrap();
271            assert_eq!(0, set.len());
272            assert!(set.is_empty());
273        });
274    }
275
276    #[test]
277    fn test_frozenset_contains() {
278        Python::attach(|py| {
279            let set = PyFrozenSet::new(py, [1]).unwrap();
280            assert!(set.contains(1).unwrap());
281        });
282    }
283
284    #[test]
285    fn test_frozenset_iter() {
286        Python::attach(|py| {
287            let set = PyFrozenSet::new(py, [1]).unwrap();
288
289            for el in set {
290                assert_eq!(1i32, el.extract::<i32>().unwrap());
291            }
292        });
293    }
294
295    #[test]
296    fn test_frozenset_iter_bound() {
297        Python::attach(|py| {
298            let set = PyFrozenSet::new(py, [1]).unwrap();
299
300            for el in &set {
301                assert_eq!(1i32, el.extract::<i32>().unwrap());
302            }
303        });
304    }
305
306    #[test]
307    fn test_frozenset_iter_size_hint() {
308        Python::attach(|py| {
309            let set = PyFrozenSet::new(py, [1]).unwrap();
310            let mut iter = set.iter();
311
312            // Exact size
313            assert_eq!(iter.len(), 1);
314            assert_eq!(iter.size_hint(), (1, Some(1)));
315            iter.next();
316            assert_eq!(iter.len(), 0);
317            assert_eq!(iter.size_hint(), (0, Some(0)));
318        });
319    }
320
321    #[test]
322    fn test_frozenset_builder() {
323        use super::PyFrozenSetBuilder;
324
325        Python::attach(|py| {
326            let mut builder = PyFrozenSetBuilder::new(py).unwrap();
327
328            // add an item
329            builder.add(1).unwrap();
330            builder.add(2).unwrap();
331            builder.add(2).unwrap();
332
333            // finalize it
334            let set = builder.finalize();
335
336            assert!(set.contains(1).unwrap());
337            assert!(set.contains(2).unwrap());
338            assert!(!set.contains(3).unwrap());
339        });
340    }
341
342    #[test]
343    fn test_iter_count() {
344        Python::attach(|py| {
345            let set = PyFrozenSet::new(py, vec![1, 2, 3]).unwrap();
346            assert_eq!(set.iter().count(), 3);
347        })
348    }
349}
pyo3/types/frozenset.rs

pyo3/types/
frozenset.rs
