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