From 5d5a9abd77496b8acecdd2c559605f4c0a5795ad Mon Sep 17 00:00:00 2001
From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>
Date: Tue, 16 May 2023 13:33:16 -0500
Subject: [PATCH] Add `PyFrozenSet.add()`

---
 newsfragments/3156.added.md |  1 +
 src/types/frozenset.rs      | 25 +++++++++++++++++++++++++
 src/types/set.rs            | 16 ++++++++--------
 3 files changed, 34 insertions(+), 8 deletions(-)
 create mode 100644 newsfragments/3156.added.md

diff --git a/newsfragments/3156.added.md b/newsfragments/3156.added.md
new file mode 100644
index 00000000000..ec4017ec22c
--- /dev/null
+++ b/newsfragments/3156.added.md
@@ -0,0 +1 @@
+Add `PyFrozenSet.add()`, which is a 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 fd2e51aa895..2a4295c0149 100644
--- a/src/types/frozenset.rs
+++ b/src/types/frozenset.rs
@@ -71,6 +71,22 @@ impl PyFrozenSet {
         }
     }
 
+    /// Adds an element to the set.
+    /// # Safety
+    ///
+    /// This is marked as unsafe because mutating a frozenset, which is conceptually immutable,
+    /// could be dangerous, even if it is allowed at the Python C ABI level.
+    /// You should only use this function if you are sure there are no other references to this frozenset.
+    pub unsafe fn add<K>(&self, key: K) -> PyResult<()>
+    where
+        K: ToPyObject,
+    {
+        err::error_on_minusone(
+            self.py(),
+            ffi::PySet_Add(self.as_ptr(), key.to_object(self.py()).as_ptr()),
+        )
+    }
+
     /// Returns an iterator of values in this frozen set.
     pub fn iter(&self) -> PyFrozenSetIterator<'_> {
         IntoIterator::into_iter(self)
@@ -219,6 +235,15 @@ mod tests {
         });
     }
 
+    #[test]
+    fn test_frozenset_add() {
+        Python::with_gil(|py| {
+            let set = PyFrozenSet::new(py, &[1, 2]).unwrap();
+            unsafe { set.add(1).unwrap() }; // Add a dupliated element
+            assert!(set.contains(1).unwrap());
+        });
+    }
+
     #[test]
     fn test_frozenset_iter() {
         Python::with_gil(|py| {
diff --git a/src/types/set.rs b/src/types/set.rs
index cf043d99dae..9f407b3c580 100644
--- a/src/types/set.rs
+++ b/src/types/set.rs
@@ -43,14 +43,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)`.
@@ -113,6 +105,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