pyo3/types/
set.rs

1use crate::types::PyIterator;
2#[cfg(feature = "gil-refs")]
3use crate::PyNativeType;
4use crate::{
5    err::{self, PyErr, PyResult},
6    ffi_ptr_ext::FfiPtrExt,
7    instance::Bound,
8    py_result_ext::PyResultExt,
9    types::any::PyAnyMethods,
10};
11use crate::{ffi, PyAny, PyObject, Python, ToPyObject};
12use std::ptr;
13
14/// Represents a Python `set`.
15///
16/// Values of this type are accessed via PyO3's smart pointers, e.g. as
17/// [`Py<PySet>`][crate::Py] or [`Bound<'py, PySet>`][Bound].
18///
19/// For APIs available on `set` objects, see the [`PySetMethods`] trait which is implemented for
20/// [`Bound<'py, PySet>`][Bound].
21#[repr(transparent)]
22pub struct PySet(PyAny);
23
24#[cfg(not(any(PyPy, GraalPy)))]
25pyobject_native_type!(
26    PySet,
27    ffi::PySetObject,
28    pyobject_native_static_type_object!(ffi::PySet_Type),
29    #checkfunction=ffi::PySet_Check
30);
31
32#[cfg(any(PyPy, GraalPy))]
33pyobject_native_type_core!(
34    PySet,
35    pyobject_native_static_type_object!(ffi::PySet_Type),
36    #checkfunction=ffi::PySet_Check
37);
38
39impl PySet {
40    /// Creates a new set with elements from the given slice.
41    ///
42    /// Returns an error if some element is not hashable.
43    #[inline]
44    pub fn new_bound<'a, 'p, T: ToPyObject + 'a>(
45        py: Python<'p>,
46        elements: impl IntoIterator<Item = &'a T>,
47    ) -> PyResult<Bound<'p, PySet>> {
48        new_from_iter(py, elements)
49    }
50
51    /// Creates a new empty set.
52    pub fn empty_bound(py: Python<'_>) -> PyResult<Bound<'_, PySet>> {
53        unsafe {
54            ffi::PySet_New(ptr::null_mut())
55                .assume_owned_or_err(py)
56                .downcast_into_unchecked()
57        }
58    }
59}
60
61#[cfg(feature = "gil-refs")]
62impl PySet {
63    /// Deprecated form of [`PySet::new_bound`].
64    #[deprecated(
65        since = "0.21.0",
66        note = "`PySet::new` will be replaced by `PySet::new_bound` in a future PyO3 version"
67    )]
68    #[inline]
69    pub fn new<'a, 'p, T: ToPyObject + 'a>(
70        py: Python<'p>,
71        elements: impl IntoIterator<Item = &'a T>,
72    ) -> PyResult<&'p PySet> {
73        Self::new_bound(py, elements).map(Bound::into_gil_ref)
74    }
75
76    /// Deprecated form of [`PySet::empty_bound`].
77    #[deprecated(
78        since = "0.21.2",
79        note = "`PySet::empty` will be replaced by `PySet::empty_bound` in a future PyO3 version"
80    )]
81    pub fn empty(py: Python<'_>) -> PyResult<&PySet> {
82        Self::empty_bound(py).map(Bound::into_gil_ref)
83    }
84
85    /// Removes all elements from the set.
86    #[inline]
87    pub fn clear(&self) {
88        self.as_borrowed().clear()
89    }
90
91    /// Returns the number of items in the set.
92    ///
93    /// This is equivalent to the Python expression `len(self)`.
94    #[inline]
95    pub fn len(&self) -> usize {
96        self.as_borrowed().len()
97    }
98
99    /// Checks if set is empty.
100    pub fn is_empty(&self) -> bool {
101        self.as_borrowed().is_empty()
102    }
103
104    /// Determines if the set contains the specified key.
105    ///
106    /// This is equivalent to the Python expression `key in self`.
107    pub fn contains<K>(&self, key: K) -> PyResult<bool>
108    where
109        K: ToPyObject,
110    {
111        self.as_borrowed().contains(key)
112    }
113
114    /// Removes the element from the set if it is present.
115    ///
116    /// Returns `true` if the element was present in the set.
117    pub fn discard<K>(&self, key: K) -> PyResult<bool>
118    where
119        K: ToPyObject,
120    {
121        self.as_borrowed().discard(key)
122    }
123
124    /// Adds an element to the set.
125    pub fn add<K>(&self, key: K) -> PyResult<()>
126    where
127        K: ToPyObject,
128    {
129        self.as_borrowed().add(key)
130    }
131
132    /// Removes and returns an arbitrary element from the set.
133    pub fn pop(&self) -> Option<PyObject> {
134        self.as_borrowed().pop().map(Bound::unbind)
135    }
136
137    /// Returns an iterator of values in this set.
138    ///
139    /// # Panics
140    ///
141    /// If PyO3 detects that the set is mutated during iteration, it will panic.
142    pub fn iter(&self) -> PySetIterator<'_> {
143        PySetIterator(BoundSetIterator::new(self.as_borrowed().to_owned()))
144    }
145}
146
147/// Implementation of functionality for [`PySet`].
148///
149/// These methods are defined for the `Bound<'py, PySet>` smart pointer, so to use method call
150/// syntax these methods are separated into a trait, because stable Rust does not yet support
151/// `arbitrary_self_types`.
152#[doc(alias = "PySet")]
153pub trait PySetMethods<'py>: crate::sealed::Sealed {
154    /// Removes all elements from the set.
155    fn clear(&self);
156
157    /// Returns the number of items in the set.
158    ///
159    /// This is equivalent to the Python expression `len(self)`.
160    fn len(&self) -> usize;
161
162    /// Checks if set is empty.
163    fn is_empty(&self) -> bool {
164        self.len() == 0
165    }
166
167    /// Determines if the set contains the specified key.
168    ///
169    /// This is equivalent to the Python expression `key in self`.
170    fn contains<K>(&self, key: K) -> PyResult<bool>
171    where
172        K: ToPyObject;
173
174    /// Removes the element from the set if it is present.
175    ///
176    /// Returns `true` if the element was present in the set.
177    fn discard<K>(&self, key: K) -> PyResult<bool>
178    where
179        K: ToPyObject;
180
181    /// Adds an element to the set.
182    fn add<K>(&self, key: K) -> PyResult<()>
183    where
184        K: ToPyObject;
185
186    /// Removes and returns an arbitrary element from the set.
187    fn pop(&self) -> Option<Bound<'py, PyAny>>;
188
189    /// Returns an iterator of values in this set.
190    ///
191    /// # Panics
192    ///
193    /// If PyO3 detects that the set is mutated during iteration, it will panic.
194    fn iter(&self) -> BoundSetIterator<'py>;
195}
196
197impl<'py> PySetMethods<'py> for Bound<'py, PySet> {
198    #[inline]
199    fn clear(&self) {
200        unsafe {
201            ffi::PySet_Clear(self.as_ptr());
202        }
203    }
204
205    #[inline]
206    fn len(&self) -> usize {
207        unsafe { ffi::PySet_Size(self.as_ptr()) as usize }
208    }
209
210    fn contains<K>(&self, key: K) -> PyResult<bool>
211    where
212        K: ToPyObject,
213    {
214        fn inner(set: &Bound<'_, PySet>, key: Bound<'_, PyAny>) -> PyResult<bool> {
215            match unsafe { ffi::PySet_Contains(set.as_ptr(), key.as_ptr()) } {
216                1 => Ok(true),
217                0 => Ok(false),
218                _ => Err(PyErr::fetch(set.py())),
219            }
220        }
221
222        let py = self.py();
223        inner(self, key.to_object(py).into_bound(py))
224    }
225
226    fn discard<K>(&self, key: K) -> PyResult<bool>
227    where
228        K: ToPyObject,
229    {
230        fn inner(set: &Bound<'_, PySet>, key: Bound<'_, PyAny>) -> PyResult<bool> {
231            match unsafe { ffi::PySet_Discard(set.as_ptr(), key.as_ptr()) } {
232                1 => Ok(true),
233                0 => Ok(false),
234                _ => Err(PyErr::fetch(set.py())),
235            }
236        }
237
238        let py = self.py();
239        inner(self, key.to_object(py).into_bound(py))
240    }
241
242    fn add<K>(&self, key: K) -> PyResult<()>
243    where
244        K: ToPyObject,
245    {
246        fn inner(set: &Bound<'_, PySet>, key: Bound<'_, PyAny>) -> PyResult<()> {
247            err::error_on_minusone(set.py(), unsafe {
248                ffi::PySet_Add(set.as_ptr(), key.as_ptr())
249            })
250        }
251
252        let py = self.py();
253        inner(self, key.to_object(py).into_bound(py))
254    }
255
256    fn pop(&self) -> Option<Bound<'py, PyAny>> {
257        let element = unsafe { ffi::PySet_Pop(self.as_ptr()).assume_owned_or_err(self.py()) };
258        match element {
259            Ok(e) => Some(e),
260            Err(_) => None,
261        }
262    }
263
264    fn iter(&self) -> BoundSetIterator<'py> {
265        BoundSetIterator::new(self.clone())
266    }
267}
268
269/// PyO3 implementation of an iterator for a Python `set` object.
270#[cfg(feature = "gil-refs")]
271pub struct PySetIterator<'py>(BoundSetIterator<'py>);
272
273#[cfg(feature = "gil-refs")]
274impl<'py> Iterator for PySetIterator<'py> {
275    type Item = &'py super::PyAny;
276
277    /// Advances the iterator and returns the next value.
278    ///
279    /// # Panics
280    ///
281    /// If PyO3 detects that the set is mutated during iteration, it will panic.
282    #[inline]
283    fn next(&mut self) -> Option<Self::Item> {
284        self.0.next().map(Bound::into_gil_ref)
285    }
286
287    fn size_hint(&self) -> (usize, Option<usize>) {
288        self.0.size_hint()
289    }
290}
291
292#[cfg(feature = "gil-refs")]
293impl ExactSizeIterator for PySetIterator<'_> {
294    fn len(&self) -> usize {
295        self.0.len()
296    }
297}
298
299#[cfg(feature = "gil-refs")]
300impl<'py> IntoIterator for &'py PySet {
301    type Item = &'py PyAny;
302    type IntoIter = PySetIterator<'py>;
303    /// Returns an iterator of values in this set.
304    ///
305    /// # Panics
306    ///
307    /// If PyO3 detects that the set is mutated during iteration, it will panic.
308    fn into_iter(self) -> Self::IntoIter {
309        PySetIterator(BoundSetIterator::new(self.as_borrowed().to_owned()))
310    }
311}
312
313impl<'py> IntoIterator for Bound<'py, PySet> {
314    type Item = Bound<'py, PyAny>;
315    type IntoIter = BoundSetIterator<'py>;
316
317    /// Returns an iterator of values in this set.
318    ///
319    /// # Panics
320    ///
321    /// If PyO3 detects that the set is mutated during iteration, it will panic.
322    fn into_iter(self) -> Self::IntoIter {
323        BoundSetIterator::new(self)
324    }
325}
326
327impl<'py> IntoIterator for &Bound<'py, PySet> {
328    type Item = Bound<'py, PyAny>;
329    type IntoIter = BoundSetIterator<'py>;
330
331    /// Returns an iterator of values in this set.
332    ///
333    /// # Panics
334    ///
335    /// If PyO3 detects that the set is mutated during iteration, it will panic.
336    fn into_iter(self) -> Self::IntoIter {
337        self.iter()
338    }
339}
340
341/// PyO3 implementation of an iterator for a Python `set` object.
342pub struct BoundSetIterator<'p> {
343    it: Bound<'p, PyIterator>,
344    // Remaining elements in the set. This is fine to store because
345    // Python will error if the set changes size during iteration.
346    remaining: usize,
347}
348
349impl<'py> BoundSetIterator<'py> {
350    pub(super) fn new(set: Bound<'py, PySet>) -> Self {
351        Self {
352            it: PyIterator::from_bound_object(&set).unwrap(),
353            remaining: set.len(),
354        }
355    }
356}
357
358impl<'py> Iterator for BoundSetIterator<'py> {
359    type Item = Bound<'py, super::PyAny>;
360
361    /// Advances the iterator and returns the next value.
362    fn next(&mut self) -> Option<Self::Item> {
363        self.remaining = self.remaining.saturating_sub(1);
364        self.it.next().map(Result::unwrap)
365    }
366
367    fn size_hint(&self) -> (usize, Option<usize>) {
368        (self.remaining, Some(self.remaining))
369    }
370}
371
372impl<'py> ExactSizeIterator for BoundSetIterator<'py> {
373    fn len(&self) -> usize {
374        self.remaining
375    }
376}
377
378#[inline]
379pub(crate) fn new_from_iter<T: ToPyObject>(
380    py: Python<'_>,
381    elements: impl IntoIterator<Item = T>,
382) -> PyResult<Bound<'_, PySet>> {
383    fn inner<'py>(
384        py: Python<'py>,
385        elements: &mut dyn Iterator<Item = PyObject>,
386    ) -> PyResult<Bound<'py, PySet>> {
387        let set = unsafe {
388            // We create the  `Py` pointer because its Drop cleans up the set if user code panics.
389            ffi::PySet_New(std::ptr::null_mut())
390                .assume_owned_or_err(py)?
391                .downcast_into_unchecked()
392        };
393        let ptr = set.as_ptr();
394
395        for obj in elements {
396            err::error_on_minusone(py, unsafe { ffi::PySet_Add(ptr, obj.as_ptr()) })?;
397        }
398
399        Ok(set)
400    }
401
402    let mut iter = elements.into_iter().map(|e| e.to_object(py));
403    inner(py, &mut iter)
404}
405
406#[cfg(test)]
407mod tests {
408    use super::PySet;
409    use crate::{
410        types::{PyAnyMethods, PySetMethods},
411        Python, ToPyObject,
412    };
413    use std::collections::HashSet;
414
415    #[test]
416    fn test_set_new() {
417        Python::with_gil(|py| {
418            let set = PySet::new_bound(py, &[1]).unwrap();
419            assert_eq!(1, set.len());
420
421            let v = vec![1];
422            assert!(PySet::new_bound(py, &[v]).is_err());
423        });
424    }
425
426    #[test]
427    fn test_set_empty() {
428        Python::with_gil(|py| {
429            let set = PySet::empty_bound(py).unwrap();
430            assert_eq!(0, set.len());
431            assert!(set.is_empty());
432        });
433    }
434
435    #[test]
436    fn test_set_len() {
437        Python::with_gil(|py| {
438            let mut v = HashSet::new();
439            let ob = v.to_object(py);
440            let set = ob.downcast_bound::<PySet>(py).unwrap();
441            assert_eq!(0, set.len());
442            v.insert(7);
443            let ob = v.to_object(py);
444            let set2 = ob.downcast_bound::<PySet>(py).unwrap();
445            assert_eq!(1, set2.len());
446        });
447    }
448
449    #[test]
450    fn test_set_clear() {
451        Python::with_gil(|py| {
452            let set = PySet::new_bound(py, &[1]).unwrap();
453            assert_eq!(1, set.len());
454            set.clear();
455            assert_eq!(0, set.len());
456        });
457    }
458
459    #[test]
460    fn test_set_contains() {
461        Python::with_gil(|py| {
462            let set = PySet::new_bound(py, &[1]).unwrap();
463            assert!(set.contains(1).unwrap());
464        });
465    }
466
467    #[test]
468    fn test_set_discard() {
469        Python::with_gil(|py| {
470            let set = PySet::new_bound(py, &[1]).unwrap();
471            assert!(!set.discard(2).unwrap());
472            assert_eq!(1, set.len());
473
474            assert!(set.discard(1).unwrap());
475            assert_eq!(0, set.len());
476            assert!(!set.discard(1).unwrap());
477
478            assert!(set.discard(vec![1, 2]).is_err());
479        });
480    }
481
482    #[test]
483    fn test_set_add() {
484        Python::with_gil(|py| {
485            let set = PySet::new_bound(py, &[1, 2]).unwrap();
486            set.add(1).unwrap(); // Add a dupliated element
487            assert!(set.contains(1).unwrap());
488        });
489    }
490
491    #[test]
492    fn test_set_pop() {
493        Python::with_gil(|py| {
494            let set = PySet::new_bound(py, &[1]).unwrap();
495            let val = set.pop();
496            assert!(val.is_some());
497            let val2 = set.pop();
498            assert!(val2.is_none());
499            assert!(py
500                .eval_bound("print('Exception state should not be set.')", None, None)
501                .is_ok());
502        });
503    }
504
505    #[test]
506    fn test_set_iter() {
507        Python::with_gil(|py| {
508            let set = PySet::new_bound(py, &[1]).unwrap();
509
510            for el in set {
511                assert_eq!(1i32, el.extract::<'_, i32>().unwrap());
512            }
513        });
514    }
515
516    #[test]
517    fn test_set_iter_bound() {
518        use crate::types::any::PyAnyMethods;
519
520        Python::with_gil(|py| {
521            let set = PySet::new_bound(py, &[1]).unwrap();
522
523            for el in &set {
524                assert_eq!(1i32, el.extract::<i32>().unwrap());
525            }
526        });
527    }
528
529    #[test]
530    #[should_panic]
531    fn test_set_iter_mutation() {
532        Python::with_gil(|py| {
533            let set = PySet::new_bound(py, &[1, 2, 3, 4, 5]).unwrap();
534
535            for _ in &set {
536                let _ = set.add(42);
537            }
538        });
539    }
540
541    #[test]
542    #[should_panic]
543    fn test_set_iter_mutation_same_len() {
544        Python::with_gil(|py| {
545            let set = PySet::new_bound(py, &[1, 2, 3, 4, 5]).unwrap();
546
547            for item in &set {
548                let item: i32 = item.extract().unwrap();
549                let _ = set.del_item(item);
550                let _ = set.add(item + 10);
551            }
552        });
553    }
554
555    #[test]
556    fn test_set_iter_size_hint() {
557        Python::with_gil(|py| {
558            let set = PySet::new_bound(py, &[1]).unwrap();
559            let mut iter = set.iter();
560
561            // Exact size
562            assert_eq!(iter.len(), 1);
563            assert_eq!(iter.size_hint(), (1, Some(1)));
564            iter.next();
565            assert_eq!(iter.len(), 0);
566            assert_eq!(iter.size_hint(), (0, Some(0)));
567        });
568    }
569}
pyo3/types/set.rs

pyo3/types/
set.rs
