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: impl Iterator> = elements.into_iter().map(|e: T| 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: Bound<'_, PySet> = unsafe {
305 // We create the `Bound` pointer because its Drop cleans up the set if
306 // user code errors or panics.
307 ffiBound<'_, PyAny>::PySet_New(arg1:std::ptr::null_mut())
308 .assume_owned_or_err(py)?
309 .downcast_into_unchecked()
310 };
311 let ptr: *mut PyObject = set.as_ptr();
312
313 elements.into_iter().try_for_each(|element: T| {
314 let obj: >::Output = element.into_pyobject_or_pyerr(py)?;
315 err::error_on_minusone(py, result:unsafe { ffi::PySet_Add(set:ptr, key: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}
499