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<'py>(Bound<'py, PyIterator>);
192
193impl<'py> BoundFrozenSetIterator<'py> {
194    pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
195        Self(PyIterator::from_object(&set).expect("frozenset should always be iterable"))
196    }
197}
198
199impl<'py> Iterator for BoundFrozenSetIterator<'py> {
200    type Item = Bound<'py, super::PyAny>;
201
202    /// Advances the iterator and returns the next value.
203    fn next(&mut self) -> Option<Self::Item> {
204        self.0
205            .next()
206            .map(|result| result.expect("frozenset iteration should be infallible"))
207    }
208
209    fn size_hint(&self) -> (usize, Option<usize>) {
210        let len = ExactSizeIterator::len(self);
211        (len, Some(len))
212    }
213
214    #[inline]
215    fn count(self) -> usize
216    where
217        Self: Sized,
218    {
219        self.len()
220    }
221}
222
223impl ExactSizeIterator for BoundFrozenSetIterator<'_> {
224    fn len(&self) -> usize {
225        self.0.size_hint().0
226    }
227}
228
229#[inline]
230pub(crate) fn try_new_from_iter<'py, T>(
231    py: Python<'py>,
232    elements: impl IntoIterator<Item = T>,
233) -> PyResult<Bound<'py, PyFrozenSet>>
234where
235    T: IntoPyObject<'py>,
236{
237    let set = unsafe {
238        // We create the  `Py` pointer because its Drop cleans up the set if user code panics.
239        ffi::PyFrozenSet_New(std::ptr::null_mut())
240            .assume_owned_or_err(py)?
241            .cast_into_unchecked()
242    };
243    let ptr = set.as_ptr();
244
245    for e in elements {
246        let obj = e.into_pyobject_or_pyerr(py)?;
247        err::error_on_minusone(py, unsafe { ffi::PySet_Add(ptr, obj.as_ptr()) })?;
248    }
249
250    Ok(set)
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256    use crate::types::PyAnyMethods as _;
257
258    #[test]
259    fn test_frozenset_new_and_len() {
260        Python::attach(|py| {
261            let set = PyFrozenSet::new(py, [1]).unwrap();
262            assert_eq!(1, set.len());
263
264            let v = vec![1];
265            assert!(PyFrozenSet::new(py, &[v]).is_err());
266        });
267    }
268
269    #[test]
270    fn test_frozenset_empty() {
271        Python::attach(|py| {
272            let set = PyFrozenSet::empty(py).unwrap();
273            assert_eq!(0, set.len());
274            assert!(set.is_empty());
275        });
276    }
277
278    #[test]
279    fn test_frozenset_contains() {
280        Python::attach(|py| {
281            let set = PyFrozenSet::new(py, [1]).unwrap();
282            assert!(set.contains(1).unwrap());
283        });
284    }
285
286    #[test]
287    fn test_frozenset_iter() {
288        Python::attach(|py| {
289            let set = PyFrozenSet::new(py, [1]).unwrap();
290
291            for el in set {
292                assert_eq!(1i32, el.extract::<i32>().unwrap());
293            }
294        });
295    }
296
297    #[test]
298    fn test_frozenset_iter_bound() {
299        Python::attach(|py| {
300            let set = PyFrozenSet::new(py, [1]).unwrap();
301
302            for el in &set {
303                assert_eq!(1i32, el.extract::<i32>().unwrap());
304            }
305        });
306    }
307
308    #[test]
309    fn test_frozenset_iter_size_hint() {
310        Python::attach(|py| {
311            let set = PyFrozenSet::new(py, [1]).unwrap();
312            let mut iter = set.iter();
313
314            // Exact size
315            assert_eq!(iter.len(), 1);
316            assert_eq!(iter.size_hint(), (1, Some(1)));
317            iter.next();
318            assert_eq!(iter.len(), 0);
319            assert_eq!(iter.size_hint(), (0, Some(0)));
320        });
321    }
322
323    #[test]
324    fn test_frozenset_builder() {
325        use super::PyFrozenSetBuilder;
326
327        Python::attach(|py| {
328            let mut builder = PyFrozenSetBuilder::new(py).unwrap();
329
330            // add an item
331            builder.add(1).unwrap();
332            builder.add(2).unwrap();
333            builder.add(2).unwrap();
334
335            // finalize it
336            let set = builder.finalize();
337
338            assert!(set.contains(1).unwrap());
339            assert!(set.contains(2).unwrap());
340            assert!(!set.contains(3).unwrap());
341        });
342    }
343
344    #[test]
345    fn test_iter_count() {
346        Python::attach(|py| {
347            let set = PyFrozenSet::new(py, vec![1, 2, 3]).unwrap();
348            assert_eq!(set.iter().count(), 3);
349        })
350    }
351}
pyo3/types/frozenset.rs

pyo3/types/
frozenset.rs
