diff --git a/newsfragments/3156.added.md b/newsfragments/3156.added.md new file mode 100644 index 00000000..75a83ef0 --- /dev/null +++ b/newsfragments/3156.added.md @@ -0,0 +1 @@ +Add `pyo3::types::PyFrozenSetBuilder` to allow building a `PyFrozenSet` item by item. \ No newline at end of file diff --git a/src/types/frozenset.rs b/src/types/frozenset.rs index fd2e51aa..ecfbbb43 100644 --- a/src/types/frozenset.rs +++ b/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> { + Ok(PyFrozenSetBuilder { + py_frozen_set: PyFrozenSet::empty(py)?, + }) + } + + /// Adds an element to the set. + pub fn add(&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,25 @@ 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(1).unwrap(); + builder.add(2).unwrap(); + builder.add(2).unwrap(); + + // finalize it + let set = builder.finalize(); + + assert!(set.contains(1).unwrap()); + assert!(set.contains(2).unwrap()); + assert!(!set.contains(3).unwrap()); + }); + } } diff --git a/src/types/mod.rs b/src/types/mod.rs index 802a20c6..06a677ab 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;