1#[cfg(Py_LIMITED_API)]
2use crate::types::PyIterator;
3use crate::{
4 err::{self, PyErr, PyResult},
5 Py, PyObject,
6};
7use crate::{ffi, PyAny, Python, ToPyObject};
8
9use std::ptr;
10
11/// Allows building a Python `frozenset` one item at a time
12pub struct PyFrozenSetBuilder<'py> {
13 py_frozen_set: &'py PyFrozenSet,
14}
15
16impl<'py> PyFrozenSetBuilder<'py> {
17 /// Create a new `FrozenSetBuilder`.
18 /// Since this allocates a `PyFrozenSet` internally it may
19 /// panic when running out of memory.
20 pub fn new(py: Python<'py>) -> PyResult<PyFrozenSetBuilder<'py>> {
21 Ok(PyFrozenSetBuilder {
22 py_frozen_set: PyFrozenSet::empty(py)?,
23 })
24 }
25
26 /// Adds an element to the set.
27 pub fn add<K>(&mut self, key: K) -> PyResult<()>
28 where
29 K: ToPyObject,
30 {
31 fn inner(frozenset: &PyFrozenSet, key: PyObject) -> PyResult<()> {
32 err::error_on_minusone(frozenset.py(), unsafe {
33 ffi::PySet_Add(frozenset.as_ptr(), key.as_ptr())
34 })
35 }
36
37 inner(self.py_frozen_set, key.to_object(self.py_frozen_set.py()))
38 }
39
40 /// Finish building the set and take ownership of its current value
41 pub fn finalize(self) -> &'py PyFrozenSet {
42 self.py_frozen_set
43 }
44}
45
46/// Represents a Python `frozenset`
47#[repr(transparent)]
48pub struct PyFrozenSet(PyAny);
49
50#[cfg(not(PyPy))]
51pyobject_native_type!(
52 PyFrozenSet,
53 ffi::PySetObject,
54 pyobject_native_static_type_object!(ffi::PyFrozenSet_Type),
55 #checkfunction=ffi::PyFrozenSet_Check
56);
57
58#[cfg(PyPy)]
59pyobject_native_type_core!(
60 PyFrozenSet,
61 pyobject_native_static_type_object!(ffi::PyFrozenSet_Type),
62 #checkfunction=ffi::PyFrozenSet_Check
63);
64
65impl PyFrozenSet {
66 /// Creates a new frozenset.
67 ///
68 /// May panic when running out of memory.
69 #[inline]
70 pub fn new<'a, 'p, T: ToPyObject + 'a>(
71 py: Python<'p>,
72 elements: impl IntoIterator<Item = &'a T>,
73 ) -> PyResult<&'p PyFrozenSet> {
74 new_from_iter(py, elements).map(|set| set.into_ref(py))
75 }
76
77 /// Creates a new empty frozen set
78 pub fn empty(py: Python<'_>) -> PyResult<&PyFrozenSet> {
79 unsafe { py.from_owned_ptr_or_err(ffi::PyFrozenSet_New(ptr::null_mut())) }
80 }
81
82 /// Return the number of items in the set.
83 /// This is equivalent to len(p) on a set.
84 #[inline]
85 pub fn len(&self) -> usize {
86 unsafe { ffi::PySet_Size(self.as_ptr()) as usize }
87 }
88
89 /// Check if set is empty.
90 pub fn is_empty(&self) -> bool {
91 self.len() == 0
92 }
93
94 /// Determine if the set contains the specified key.
95 /// This is equivalent to the Python expression `key in self`.
96 pub fn contains<K>(&self, key: K) -> PyResult<bool>
97 where
98 K: ToPyObject,
99 {
100 fn inner(frozenset: &PyFrozenSet, key: PyObject) -> PyResult<bool> {
101 match unsafe { ffi::PySet_Contains(frozenset.as_ptr(), key.as_ptr()) } {
102 1 => Ok(true),
103 0 => Ok(false),
104 _ => Err(PyErr::fetch(frozenset.py())),
105 }
106 }
107
108 inner(self, key.to_object(self.py()))
109 }
110
111 /// Returns an iterator of values in this frozen set.
112 pub fn iter(&self) -> PyFrozenSetIterator<'_> {
113 IntoIterator::into_iter(self)
114 }
115}
116
117#[cfg(Py_LIMITED_API)]
118mod impl_ {
119 use super::*;
120
121 impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
122 type Item = &'a PyAny;
123 type IntoIter = PyFrozenSetIterator<'a>;
124
125 fn into_iter(self) -> Self::IntoIter {
126 PyFrozenSetIterator {
127 it: PyIterator::from_object(self).unwrap(),
128 }
129 }
130 }
131
132 /// PyO3 implementation of an iterator for a Python `frozenset` object.
133 pub struct PyFrozenSetIterator<'p> {
134 it: &'p PyIterator,
135 }
136
137 impl<'py> Iterator for PyFrozenSetIterator<'py> {
138 type Item = &'py super::PyAny;
139
140 #[inline]
141 fn next(&mut self) -> Option<Self::Item> {
142 self.it.next().map(Result::unwrap)
143 }
144 }
145}
146
147#[cfg(not(Py_LIMITED_API))]
148mod impl_ {
149 use super::*;
150
151 impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
152 type Item = &'a PyAny;
153 type IntoIter = PyFrozenSetIterator<'a>;
154
155 fn into_iter(self) -> Self::IntoIter {
156 PyFrozenSetIterator { set: self, pos: 0 }
157 }
158 }
159
160 /// PyO3 implementation of an iterator for a Python `frozenset` object.
161 pub struct PyFrozenSetIterator<'py> {
162 set: &'py PyFrozenSet,
163 pos: ffi::Py_ssize_t,
164 }
165
166 impl<'py> Iterator for PyFrozenSetIterator<'py> {
167 type Item = &'py PyAny;
168
169 #[inline]
170 fn next(&mut self) -> Option<Self::Item> {
171 unsafe {
172 let mut key: *mut ffi::PyObject = std::ptr::null_mut();
173 let mut hash: ffi::Py_hash_t = 0;
174 if ffi::_PySet_NextEntry(self.set.as_ptr(), &mut self.pos, &mut key, &mut hash) != 0
175 {
176 // _PySet_NextEntry returns borrowed object; for safety must make owned (see #890)
177 Some(self.set.py().from_owned_ptr(ffi::_Py_NewRef(key)))
178 } else {
179 None
180 }
181 }
182 }
183
184 #[inline]
185 fn size_hint(&self) -> (usize, Option<usize>) {
186 let len = self.len();
187 (len, Some(len))
188 }
189 }
190
191 impl<'py> ExactSizeIterator for PyFrozenSetIterator<'py> {
192 fn len(&self) -> usize {
193 self.set.len().saturating_sub(self.pos as usize)
194 }
195 }
196}
197
198pub use impl_::*;
199
200#[inline]
201pub(crate) fn new_from_iter<T: ToPyObject>(
202 py: Python<'_>,
203 elements: impl IntoIterator<Item = T>,
204) -> PyResult<Py<PyFrozenSet>> {
205 fn inner(
206 py: Python<'_>,
207 elements: &mut dyn Iterator<Item = PyObject>,
208 ) -> PyResult<Py<PyFrozenSet>> {
209 let set: Py<PyFrozenSet> = unsafe {
210 // We create the `Py` pointer because its Drop cleans up the set if user code panics.
211 Py::from_owned_ptr_or_err(py, ptr:ffi::PyFrozenSet_New(arg1:std::ptr::null_mut()))?
212 };
213 let ptr: *mut PyObject = set.as_ptr();
214
215 for obj: Py in elements {
216 err::error_on_minusone(py, result:unsafe { ffi::PySet_Add(set:ptr, key:obj.as_ptr()) })?;
217 }
218
219 Ok(set)
220 }
221
222 let mut iter: impl Iterator> = elements.into_iter().map(|e: T| e.to_object(py));
223 inner(py, &mut iter)
224}
225
226#[cfg(test)]
227mod tests {
228 use super::*;
229
230 #[test]
231 fn test_frozenset_new_and_len() {
232 Python::with_gil(|py| {
233 let set = PyFrozenSet::new(py, &[1]).unwrap();
234 assert_eq!(1, set.len());
235
236 let v = vec![1];
237 assert!(PyFrozenSet::new(py, &[v]).is_err());
238 });
239 }
240
241 #[test]
242 fn test_frozenset_empty() {
243 Python::with_gil(|py| {
244 let set = PyFrozenSet::empty(py).unwrap();
245 assert_eq!(0, set.len());
246 });
247 }
248
249 #[test]
250 fn test_frozenset_contains() {
251 Python::with_gil(|py| {
252 let set = PyFrozenSet::new(py, &[1]).unwrap();
253 assert!(set.contains(1).unwrap());
254 });
255 }
256
257 #[test]
258 fn test_frozenset_iter() {
259 Python::with_gil(|py| {
260 let set = PyFrozenSet::new(py, &[1]).unwrap();
261
262 // iter method
263 for el in set {
264 assert_eq!(1i32, el.extract::<i32>().unwrap());
265 }
266
267 // intoiterator iteration
268 for el in set {
269 assert_eq!(1i32, el.extract::<i32>().unwrap());
270 }
271 });
272 }
273
274 #[test]
275 fn test_frozenset_builder() {
276 use super::PyFrozenSetBuilder;
277
278 Python::with_gil(|py| {
279 let mut builder = PyFrozenSetBuilder::new(py).unwrap();
280
281 // add an item
282 builder.add(1).unwrap();
283 builder.add(2).unwrap();
284 builder.add(2).unwrap();
285
286 // finalize it
287 let set = builder.finalize();
288
289 assert!(set.contains(1).unwrap());
290 assert!(set.contains(2).unwrap());
291 assert!(!set.contains(3).unwrap());
292 });
293 }
294}
295