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

pyo3/types/
frozenset.rs
