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

pyo3/types/
frozenset.rs
