pyo3/types/
set.rs

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

pyo3/types/
set.rs
