pyo3/types/
set.rs

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

pyo3/types/
set.rs
