use crate::types::PyIterator;
use crate::{
err::{self, PyErr, PyResult},
ffi_ptr_ext::FfiPtrExt,
instance::Bound,
py_result_ext::PyResultExt,
types::any::PyAnyMethods,
PyNativeType,
};
use crate::{ffi, PyAny, PyObject, Python, ToPyObject};
use std::ptr;
#[repr(transparent)]
pub struct PySet(PyAny);
#[cfg(not(any(PyPy, GraalPy)))]
pyobject_native_type!(
PySet,
ffi::PySetObject,
pyobject_native_static_type_object!(ffi::PySet_Type),
#checkfunction=ffi::PySet_Check
);
#[cfg(any(PyPy, GraalPy))]
pyobject_native_type_core!(
PySet,
pyobject_native_static_type_object!(ffi::PySet_Type),
#checkfunction=ffi::PySet_Check
);
impl PySet {
#[cfg_attr(
not(feature = "gil-refs"),
deprecated(
since = "0.21.0",
note = "`PySet::new` will be replaced by `PySet::new_bound` in a future PyO3 version"
)
)]
#[inline]
pub fn new<'a, 'p, T: ToPyObject + 'a>(
py: Python<'p>,
elements: impl IntoIterator<Item = &'a T>,
) -> PyResult<&'p PySet> {
Self::new_bound(py, elements).map(Bound::into_gil_ref)
}
#[inline]
pub fn new_bound<'a, 'p, T: ToPyObject + 'a>(
py: Python<'p>,
elements: impl IntoIterator<Item = &'a T>,
) -> PyResult<Bound<'p, PySet>> {
new_from_iter(py, elements)
}
#[cfg_attr(
not(feature = "gil-refs"),
deprecated(
since = "0.21.2",
note = "`PySet::empty` will be replaced by `PySet::empty_bound` in a future PyO3 version"
)
)]
pub fn empty(py: Python<'_>) -> PyResult<&PySet> {
Self::empty_bound(py).map(Bound::into_gil_ref)
}
pub fn empty_bound(py: Python<'_>) -> PyResult<Bound<'_, PySet>> {
unsafe {
ffi::PySet_New(ptr::null_mut())
.assume_owned_or_err(py)
.downcast_into_unchecked()
}
}
#[inline]
pub fn clear(&self) {
self.as_borrowed().clear()
}
#[inline]
pub fn len(&self) -> usize {
self.as_borrowed().len()
}
pub fn is_empty(&self) -> bool {
self.as_borrowed().is_empty()
}
pub fn contains<K>(&self, key: K) -> PyResult<bool>
where
K: ToPyObject,
{
self.as_borrowed().contains(key)
}
pub fn discard<K>(&self, key: K) -> PyResult<bool>
where
K: ToPyObject,
{
self.as_borrowed().discard(key)
}
pub fn add<K>(&self, key: K) -> PyResult<()>
where
K: ToPyObject,
{
self.as_borrowed().add(key)
}
pub fn pop(&self) -> Option<PyObject> {
self.as_borrowed().pop().map(Bound::unbind)
}
pub fn iter(&self) -> PySetIterator<'_> {
PySetIterator(BoundSetIterator::new(self.as_borrowed().to_owned()))
}
}
#[doc(alias = "PySet")]
pub trait PySetMethods<'py>: crate::sealed::Sealed {
fn clear(&self);
fn len(&self) -> usize;
fn is_empty(&self) -> bool {
self.len() == 0
}
fn contains<K>(&self, key: K) -> PyResult<bool>
where
K: ToPyObject;
fn discard<K>(&self, key: K) -> PyResult<bool>
where
K: ToPyObject;
fn add<K>(&self, key: K) -> PyResult<()>
where
K: ToPyObject;
fn pop(&self) -> Option<Bound<'py, PyAny>>;
fn iter(&self) -> BoundSetIterator<'py>;
}
impl<'py> PySetMethods<'py> for Bound<'py, PySet> {
#[inline]
fn clear(&self) {
unsafe {
ffi::PySet_Clear(self.as_ptr());
}
}
#[inline]
fn len(&self) -> usize {
unsafe { ffi::PySet_Size(self.as_ptr()) as usize }
}
fn contains<K>(&self, key: K) -> PyResult<bool>
where
K: ToPyObject,
{
fn inner(set: &Bound<'_, PySet>, key: Bound<'_, PyAny>) -> PyResult<bool> {
match unsafe { ffi::PySet_Contains(set.as_ptr(), key.as_ptr()) } {
1 => Ok(true),
0 => Ok(false),
_ => Err(PyErr::fetch(set.py())),
}
}
let py = self.py();
inner(self, key.to_object(py).into_bound(py))
}
fn discard<K>(&self, key: K) -> PyResult<bool>
where
K: ToPyObject,
{
fn inner(set: &Bound<'_, PySet>, key: Bound<'_, PyAny>) -> PyResult<bool> {
match unsafe { ffi::PySet_Discard(set.as_ptr(), key.as_ptr()) } {
1 => Ok(true),
0 => Ok(false),
_ => Err(PyErr::fetch(set.py())),
}
}
let py = self.py();
inner(self, key.to_object(py).into_bound(py))
}
fn add<K>(&self, key: K) -> PyResult<()>
where
K: ToPyObject,
{
fn inner(set: &Bound<'_, PySet>, key: Bound<'_, PyAny>) -> PyResult<()> {
err::error_on_minusone(set.py(), unsafe {
ffi::PySet_Add(set.as_ptr(), key.as_ptr())
})
}
let py = self.py();
inner(self, key.to_object(py).into_bound(py))
}
fn pop(&self) -> Option<Bound<'py, PyAny>> {
let element = unsafe { ffi::PySet_Pop(self.as_ptr()).assume_owned_or_err(self.py()) };
match element {
Ok(e) => Some(e),
Err(_) => None,
}
}
fn iter(&self) -> BoundSetIterator<'py> {
BoundSetIterator::new(self.clone())
}
}
pub struct PySetIterator<'py>(BoundSetIterator<'py>);
impl<'py> Iterator for PySetIterator<'py> {
type Item = &'py super::PyAny;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(Bound::into_gil_ref)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
impl ExactSizeIterator for PySetIterator<'_> {
fn len(&self) -> usize {
self.0.len()
}
}
impl<'py> IntoIterator for &'py PySet {
type Item = &'py PyAny;
type IntoIter = PySetIterator<'py>;
fn into_iter(self) -> Self::IntoIter {
PySetIterator(BoundSetIterator::new(self.as_borrowed().to_owned()))
}
}
impl<'py> IntoIterator for Bound<'py, PySet> {
type Item = Bound<'py, PyAny>;
type IntoIter = BoundSetIterator<'py>;
fn into_iter(self) -> Self::IntoIter {
BoundSetIterator::new(self)
}
}
impl<'py> IntoIterator for &Bound<'py, PySet> {
type Item = Bound<'py, PyAny>;
type IntoIter = BoundSetIterator<'py>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
pub struct BoundSetIterator<'p> {
it: Bound<'p, PyIterator>,
remaining: usize,
}
impl<'py> BoundSetIterator<'py> {
pub(super) fn new(set: Bound<'py, PySet>) -> Self {
Self {
it: PyIterator::from_bound_object(&set).unwrap(),
remaining: set.len(),
}
}
}
impl<'py> Iterator for BoundSetIterator<'py> {
type Item = Bound<'py, super::PyAny>;
fn next(&mut self) -> Option<Self::Item> {
self.remaining = self.remaining.saturating_sub(1);
self.it.next().map(Result::unwrap)
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.remaining, Some(self.remaining))
}
}
impl<'py> ExactSizeIterator for BoundSetIterator<'py> {
fn len(&self) -> usize {
self.remaining
}
}
#[inline]
pub(crate) fn new_from_iter<T: ToPyObject>(
py: Python<'_>,
elements: impl IntoIterator<Item = T>,
) -> PyResult<Bound<'_, PySet>> {
fn inner<'py>(
py: Python<'py>,
elements: &mut dyn Iterator<Item = PyObject>,
) -> PyResult<Bound<'py, PySet>> {
let set = unsafe {
ffi::PySet_New(std::ptr::null_mut())
.assume_owned_or_err(py)?
.downcast_into_unchecked()
};
let ptr = set.as_ptr();
for obj in elements {
err::error_on_minusone(py, unsafe { ffi::PySet_Add(ptr, obj.as_ptr()) })?;
}
Ok(set)
}
let mut iter = elements.into_iter().map(|e| e.to_object(py));
inner(py, &mut iter)
}
#[cfg(test)]
#[cfg_attr(not(feature = "gil-refs"), allow(deprecated))]
mod tests {
use super::PySet;
use crate::{Python, ToPyObject};
use std::collections::HashSet;
#[test]
fn test_set_new() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1]).unwrap();
assert_eq!(1, set.len());
let v = vec![1];
assert!(PySet::new(py, &[v]).is_err());
});
}
#[test]
fn test_set_empty() {
Python::with_gil(|py| {
let set = PySet::empty(py).unwrap();
assert_eq!(0, set.len());
assert!(set.is_empty());
});
}
#[test]
fn test_set_len() {
Python::with_gil(|py| {
let mut v = HashSet::new();
let ob = v.to_object(py);
let set: &PySet = ob.downcast(py).unwrap();
assert_eq!(0, set.len());
v.insert(7);
let ob = v.to_object(py);
let set2: &PySet = ob.downcast(py).unwrap();
assert_eq!(1, set2.len());
});
}
#[test]
fn test_set_clear() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1]).unwrap();
assert_eq!(1, set.len());
set.clear();
assert_eq!(0, set.len());
});
}
#[test]
fn test_set_contains() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1]).unwrap();
assert!(set.contains(1).unwrap());
});
}
#[test]
fn test_set_discard() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1]).unwrap();
assert!(!set.discard(2).unwrap());
assert_eq!(1, set.len());
assert!(set.discard(1).unwrap());
assert_eq!(0, set.len());
assert!(!set.discard(1).unwrap());
assert!(set.discard(vec![1, 2]).is_err());
});
}
#[test]
fn test_set_add() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1, 2]).unwrap();
set.add(1).unwrap(); assert!(set.contains(1).unwrap());
});
}
#[test]
fn test_set_pop() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1]).unwrap();
let val = set.pop();
assert!(val.is_some());
let val2 = set.pop();
assert!(val2.is_none());
assert!(py
.eval("print('Exception state should not be set.')", None, None)
.is_ok());
});
}
#[test]
fn test_set_iter() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1]).unwrap();
for el in set {
assert_eq!(1i32, el.extract::<'_, i32>().unwrap());
}
});
}
#[test]
fn test_set_iter_bound() {
use crate::types::any::PyAnyMethods;
Python::with_gil(|py| {
let set = PySet::new_bound(py, &[1]).unwrap();
for el in &set {
assert_eq!(1i32, el.extract::<i32>().unwrap());
}
});
}
#[test]
#[should_panic]
fn test_set_iter_mutation() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1, 2, 3, 4, 5]).unwrap();
for _ in set {
let _ = set.add(42);
}
});
}
#[test]
#[should_panic]
fn test_set_iter_mutation_same_len() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1, 2, 3, 4, 5]).unwrap();
for item in set {
let item: i32 = item.extract().unwrap();
let _ = set.del_item(item);
let _ = set.add(item + 10);
}
});
}
#[test]
fn test_set_iter_size_hint() {
Python::with_gil(|py| {
let set = PySet::new(py, &[1]).unwrap();
let mut iter = set.iter();
assert_eq!(iter.len(), 1);
assert_eq!(iter.size_hint(), (1, Some(1)));
iter.next();
assert_eq!(iter.len(), 0);
assert_eq!(iter.size_hint(), (0, Some(0)));
});
}
}