Add PyFrozenSetBuilder

newsfragments/3156.added.md

@@ -0,0 +1 @@
+Add `pyo3::types::PyFrozenSetBuilder` to allow building a `PyFrozenSet` item by item.

src/types/frozenset.rs

@@ -8,6 +8,38 @@ use crate::{ffi, AsPyPointer, PyAny, Python, ToPyObject};
 
 use std::ptr;
 
+/// Allows building a Python `frozenset` one item at a time
+pub struct PyFrozenSetBuilder<'py> {
+    py_frozen_set: &'py PyFrozenSet,
+}
+
+impl<'py> PyFrozenSetBuilder<'py> {
+    /// Create a new `FrozenSetBuilder`.
+    /// Since this allocates a `PyFrozenSet` internally it may
+    /// panic when running out of memory.
+    pub fn new(py: Python<'py>) -> PyResult<PyFrozenSetBuilder<'py>> {
+        Ok(PyFrozenSetBuilder {
+            py_frozen_set: PyFrozenSet::empty(py)?,
+        })
+    }
+
+    /// Adds an element to the set.
+    pub fn add<K>(&mut self, key: K) -> PyResult<()>
+    where
+        K: ToPyObject,
+    {
+        let py = self.py_frozen_set.py();
+        err::error_on_minusone(py, unsafe {
+            ffi::PySet_Add(self.py_frozen_set.as_ptr(), key.to_object(py).as_ptr())
+        })
+    }
+
+    /// Finish building the set and take ownership of its current value
+    pub fn finalize(self) -> &'py PyFrozenSet {
+        self.py_frozen_set
+    }
+}
+
 /// Represents a  Python `frozenset`
 #[repr(transparent)]
 pub struct PyFrozenSet(PyAny);
@@ -235,4 +267,26 @@ mod tests {
             }
         });
     }
+
+    #[test]
+    fn test_frozenset_builder() {
+        use super::PyFrozenSetBuilder;
+
+        Python::with_gil(|py| {
+            let mut builder = PyFrozenSetBuilder::new(py).unwrap();
+
+            // add an item
+            builder.add(2).unwrap();
+            builder.add(2).unwrap();
+            builder.add(3).unwrap();
+
+            // finalize it
+            let set = builder.finalize();
+
+            assert!(set.contains(1).unwrap());
+            assert!(set.contains(2).unwrap());
+            assert!(set.contains(3).unwrap());
+            assert!(!set.contains(4).unwrap());
+        });
+    }
 }

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;