From bdee4426d5366bdfe5222c2b2bc92704d5c0d0c0 Mon Sep 17 00:00:00 2001
From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>
Date: Tue, 13 Jun 2023 00:03:52 -0500
Subject: [PATCH] Add builder

---
 newsfragments/3156.added.md |  3 +-
 src/types/frozenset.rs      | 57 +++++++++++++++++++++++++++++++++++++
 src/types/mod.rs            |  2 +-
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/newsfragments/3156.added.md b/newsfragments/3156.added.md
index ec4017ec22c..f3bf3c29776 100644
--- a/newsfragments/3156.added.md
+++ b/newsfragments/3156.added.md
@@ -1 +1,2 @@
-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)
+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).
+This method is unsafe so a PyFrozenSetBuilder was also added that wraps this functionality in safe Rust.
\ No newline at end of file
diff --git a/src/types/frozenset.rs b/src/types/frozenset.rs
index 2a4295c0149..5818a8258ce 100644
--- a/src/types/frozenset.rs
+++ b/src/types/frozenset.rs
@@ -8,6 +8,41 @@ use crate::{ffi, AsPyPointer, PyAny, Python, ToPyObject};
 
 use std::ptr;
 
+/// Allows building a Python `frozenset` one item at a time
+pub struct PyFrozenSetBuilder {
+    py_frozen_set: Py<PyFrozenSet>,
+}
+
+impl PyFrozenSetBuilder {
+    /// Create a new `FrozenSetBuilder`.
+    /// Since this allocates a `PyFrozenSet` internally it may
+    /// panic when running out of memory.
+    pub fn new<'a, 'p, T: ToPyObject + 'a>(
+        py: Python<'p>,
+        elements: impl IntoIterator<Item = &'a T>,
+    ) -> PyResult<PyFrozenSetBuilder> {
+        Ok(PyFrozenSetBuilder {
+            py_frozen_set: new_from_iter(py, elements)?,
+        })
+    }
+
+    /// Adds an element to the set.
+    pub fn add<K>(&mut self, py: Python<'_>, key: K) -> PyResult<()>
+    where
+        K: ToPyObject,
+    {
+        let inner = self.py_frozen_set.as_ref(py);
+        err::error_on_minusone(py, unsafe {
+            ffi::PySet_Add(inner.as_ptr(), key.to_object(py).as_ptr())
+        })
+    }
+
+    /// Finish building the set and take ownership of its current value
+    pub fn finalize(self, py: Python<'_>) -> &PyFrozenSet {
+        self.py_frozen_set.into_ref(py)
+    }
+}
+
 /// Represents a  Python `frozenset`
 #[repr(transparent)]
 pub struct PyFrozenSet(PyAny);
@@ -260,4 +295,26 @@ mod tests {
             }
         });
     }
+
+    #[test]
+    fn test_frozenset_builder() {
+        use super::PyFrozenSetBuilder;
+
+        Python::with_gil(|py| {
+            let mut builder = PyFrozenSetBuilder::new(py, &[1]).unwrap();
+
+            // add an item
+            builder.add(py, 2).unwrap();
+            builder.add(py, 2).unwrap();
+            builder.add(py, 3).unwrap();
+
+            // finalize it
+            let set = builder.finalize(py);
+
+            assert!(set.contains(1).unwrap());
+            assert!(set.contains(2).unwrap());
+            assert!(set.contains(3).unwrap());
+            assert!(!set.contains(4).unwrap());
+        });
+    }
 }
diff --git a/src/types/mod.rs b/src/types/mod.rs
index 802a20c6f4a..06a677ab62d 100644
--- a/src/types/mod.rs
+++ b/src/types/mod.rs
@@ -19,7 +19,7 @@ pub use self::dict::{PyDictItems, PyDictKeys, PyDictValues};
 pub use self::floatob::PyFloat;
 #[cfg(all(not(Py_LIMITED_API), not(PyPy)))]
 pub use self::frame::PyFrame;
-pub use self::frozenset::PyFrozenSet;
+pub use self::frozenset::{PyFrozenSet, PyFrozenSetBuilder};
 pub use self::function::PyCFunction;
 #[cfg(all(not(Py_LIMITED_API), not(PyPy)))]
 pub use self::function::PyFunction;