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