diff --git a/newsfragments/3156.added.md b/newsfragments/3156.added.md new file mode 100644 index 00000000000..346f5d1e4dc --- /dev/null +++ b/newsfragments/3156.added.md @@ -0,0 +1 @@ +Add `clear()`, `add()`, `pop()` and `discard()` for frozenset, which are all valid operations at the C ABI level according to the [C API docs for sets](https://docs.python.org/3/c-api/set.html) diff --git a/src/types/frozenset.rs b/src/types/frozenset.rs index 5b728784a30..99d2c933d27 100644 --- a/src/types/frozenset.rs +++ b/src/types/frozenset.rs @@ -74,6 +74,47 @@ impl PyFrozenSet { } } + /// Removes the element from the set if it is present. + pub fn discard(&self, key: K) + where + K: ToPyObject, + { + unsafe { + ffi::PySet_Discard(self.as_ptr(), key.to_object(self.py()).as_ptr()); + } + } + + /// Adds an element to the set. + pub fn add(&self, key: K) -> PyResult<()> + where + K: ToPyObject, + { + unsafe { + err::error_on_minusone( + self.py(), + ffi::PySet_Add(self.as_ptr(), key.to_object(self.py()).as_ptr()), + ) + } + } + + /// Removes and returns an arbitrary element from the set. + pub fn pop(&self) -> Option { + let element = + unsafe { PyObject::from_owned_ptr_or_err(self.py(), ffi::PySet_Pop(self.as_ptr())) }; + match element { + Ok(e) => Some(e), + Err(_) => None, + } + } + + /// Removes all elements from the set. + #[inline] + pub fn clear(&self) { + unsafe { + ffi::PySet_Clear(self.as_ptr()); + } + } + /// Returns an iterator of values in this frozen set. pub fn iter(&self) -> PyFrozenSetIterator<'_> { IntoIterator::into_iter(self) @@ -206,6 +247,16 @@ mod tests { }); } + #[test] + fn test_frozenset_clear() { + Python::with_gil(|py| { + let set = PyFrozenSet::new(py, &[1]).unwrap(); + assert_eq!(1, set.len()); + set.clear(); + assert_eq!(0, set.len()); + }); + } + #[test] fn test_frozenset_empty() { Python::with_gil(|py| { @@ -222,6 +273,40 @@ mod tests { }); } + #[test] + fn test_frozenset_discard() { + Python::with_gil(|py| { + let set = PyFrozenSet::new(py, &[1]).unwrap(); + set.discard(2); + assert_eq!(1, set.len()); + set.discard(1); + assert_eq!(0, set.len()); + }); + } + + #[test] + fn test_frozenset_add() { + Python::with_gil(|py| { + let set = PyFrozenSet::new(py, &[1, 2]).unwrap(); + set.add(1).unwrap(); // Add a dupliated element + assert!(set.contains(1).unwrap()); + }); + } + + #[test] + fn test_frozenset_pop() { + Python::with_gil(|py| { + let set = PyFrozenSet::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_frozenset_iter() { Python::with_gil(|py| { diff --git a/src/types/set.rs b/src/types/set.rs index 36774b9a014..1f67b74bc83 100644 --- a/src/types/set.rs +++ b/src/types/set.rs @@ -46,14 +46,6 @@ impl PySet { unsafe { py.from_owned_ptr_or_err(ffi::PySet_New(ptr::null_mut())) } } - /// Removes all elements from the set. - #[inline] - pub fn clear(&self) { - unsafe { - ffi::PySet_Clear(self.as_ptr()); - } - } - /// Returns the number of items in the set. /// /// This is equivalent to the Python expression `len(self)`. @@ -116,6 +108,14 @@ impl PySet { } } + /// Removes all elements from the set. + #[inline] + pub fn clear(&self) { + unsafe { + ffi::PySet_Clear(self.as_ptr()); + } + } + /// Returns an iterator of values in this set. /// /// # Panics