pyo3/types/
set.rs

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

pyo3/types/
set.rs
