Protect iterators against concurrent modification (#2380)
This commit is contained in:
parent
f84c740e41
commit
4f9d3d7306
|
@ -26,6 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||||
- Allow `#[classattr]` methods to be fallible. [#2385](https://github.com/PyO3/pyo3/pull/2385)
|
- Allow `#[classattr]` methods to be fallible. [#2385](https://github.com/PyO3/pyo3/pull/2385)
|
||||||
- Prevent multiple `#[pymethods]` with the same name for a single `#[pyclass]`. [#2399](https://github.com/PyO3/pyo3/pull/2399)
|
- Prevent multiple `#[pymethods]` with the same name for a single `#[pyclass]`. [#2399](https://github.com/PyO3/pyo3/pull/2399)
|
||||||
- Fixup `lib_name` when using `PYO3_CONFIG_FILE`. [#2404](https://github.com/PyO3/pyo3/pull/2404)
|
- Fixup `lib_name` when using `PYO3_CONFIG_FILE`. [#2404](https://github.com/PyO3/pyo3/pull/2404)
|
||||||
|
- Iterators over `PySet` and `PyDict` will now panic if the underlying collection is mutated during the iteration. [#2380](https://github.com/PyO3/pyo3/pull/2380)
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
||||||
|
|
|
@ -1,6 +1,8 @@
|
||||||
// Copyright (c) 2017-present PyO3 Project and Contributors
|
// Copyright (c) 2017-present PyO3 Project and Contributors
|
||||||
|
|
||||||
|
use super::PyMapping;
|
||||||
use crate::err::{self, PyErr, PyResult};
|
use crate::err::{self, PyErr, PyResult};
|
||||||
|
use crate::ffi::Py_ssize_t;
|
||||||
use crate::types::{PyAny, PyList};
|
use crate::types::{PyAny, PyList};
|
||||||
#[cfg(not(PyPy))]
|
#[cfg(not(PyPy))]
|
||||||
use crate::IntoPyPointer;
|
use crate::IntoPyPointer;
|
||||||
|
@ -9,8 +11,6 @@ use std::collections::{BTreeMap, HashMap};
|
||||||
use std::ptr::NonNull;
|
use std::ptr::NonNull;
|
||||||
use std::{cmp, collections, hash};
|
use std::{cmp, collections, hash};
|
||||||
|
|
||||||
use super::PyMapping;
|
|
||||||
|
|
||||||
/// Represents a Python `dict`.
|
/// Represents a Python `dict`.
|
||||||
#[repr(transparent)]
|
#[repr(transparent)]
|
||||||
pub struct PyDict(PyAny);
|
pub struct PyDict(PyAny);
|
||||||
|
@ -102,7 +102,19 @@ impl PyDict {
|
||||||
///
|
///
|
||||||
/// This is equivalent to the Python expression `len(self)`.
|
/// This is equivalent to the Python expression `len(self)`.
|
||||||
pub fn len(&self) -> usize {
|
pub fn len(&self) -> usize {
|
||||||
unsafe { ffi::PyDict_Size(self.as_ptr()) as usize }
|
self._len() as usize
|
||||||
|
}
|
||||||
|
|
||||||
|
fn _len(&self) -> Py_ssize_t {
|
||||||
|
#[cfg(any(not(Py_3_8), PyPy, Py_LIMITED_API))]
|
||||||
|
unsafe {
|
||||||
|
ffi::PyDict_Size(self.as_ptr())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(all(Py_3_8, not(PyPy), not(Py_LIMITED_API)))]
|
||||||
|
unsafe {
|
||||||
|
(*self.as_ptr().cast::<ffi::PyDictObject>()).ma_used
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Checks if the dict is empty, i.e. `len(self) == 0`.
|
/// Checks if the dict is empty, i.e. `len(self) == 0`.
|
||||||
|
@ -213,13 +225,13 @@ impl PyDict {
|
||||||
|
|
||||||
/// Returns an iterator of `(key, value)` pairs in this dictionary.
|
/// Returns an iterator of `(key, value)` pairs in this dictionary.
|
||||||
///
|
///
|
||||||
/// Note that it's unsafe to use when the dictionary might be changed by
|
/// # Panics
|
||||||
/// other code.
|
///
|
||||||
|
/// If PyO3 detects that the dictionary is mutated during iteration, it will panic.
|
||||||
|
/// It is allowed to modify values as you iterate over the dictionary, but only
|
||||||
|
/// so long as the set of keys does not change.
|
||||||
pub fn iter(&self) -> PyDictIterator<'_> {
|
pub fn iter(&self) -> PyDictIterator<'_> {
|
||||||
PyDictIterator {
|
IntoIterator::into_iter(self)
|
||||||
dict: self.as_ref(),
|
|
||||||
pos: 0,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns `self` cast as a `PyMapping`.
|
/// Returns `self` cast as a `PyMapping`.
|
||||||
|
@ -229,8 +241,10 @@ impl PyDict {
|
||||||
}
|
}
|
||||||
|
|
||||||
pub struct PyDictIterator<'py> {
|
pub struct PyDictIterator<'py> {
|
||||||
dict: &'py PyAny,
|
dict: &'py PyDict,
|
||||||
pos: isize,
|
ppos: ffi::Py_ssize_t,
|
||||||
|
di_used: ffi::Py_ssize_t,
|
||||||
|
len: ffi::Py_ssize_t,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<'py> Iterator for PyDictIterator<'py> {
|
impl<'py> Iterator for PyDictIterator<'py> {
|
||||||
|
@ -238,29 +252,43 @@ impl<'py> Iterator for PyDictIterator<'py> {
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
fn next(&mut self) -> Option<Self::Item> {
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
unsafe {
|
let ma_used = self.dict._len();
|
||||||
let mut key: *mut ffi::PyObject = std::ptr::null_mut();
|
|
||||||
let mut value: *mut ffi::PyObject = std::ptr::null_mut();
|
// These checks are similar to what CPython does.
|
||||||
if ffi::PyDict_Next(self.dict.as_ptr(), &mut self.pos, &mut key, &mut value) != 0 {
|
//
|
||||||
let py = self.dict.py();
|
// If the dimension of the dict changes e.g. key-value pairs are removed
|
||||||
// PyDict_Next returns borrowed values; for safety must make them owned (see #890)
|
// or added during iteration, this will panic next time when `next` is called
|
||||||
Some((
|
if self.di_used != ma_used {
|
||||||
py.from_owned_ptr(ffi::_Py_NewRef(key)),
|
self.di_used = -1;
|
||||||
py.from_owned_ptr(ffi::_Py_NewRef(value)),
|
panic!("dictionary changed size during iteration");
|
||||||
))
|
};
|
||||||
} else {
|
|
||||||
None
|
// If the dict is changed in such a way that the length remains constant
|
||||||
}
|
// then this will panic at the end of iteration - similar to this:
|
||||||
|
//
|
||||||
|
// d = {"a":1, "b":2, "c": 3}
|
||||||
|
//
|
||||||
|
// for k, v in d.items():
|
||||||
|
// d[f"{k}_"] = 4
|
||||||
|
// del d[k]
|
||||||
|
// print(k)
|
||||||
|
//
|
||||||
|
if self.len == -1 {
|
||||||
|
self.di_used = -1;
|
||||||
|
panic!("dictionary keys changed during iteration");
|
||||||
|
};
|
||||||
|
|
||||||
|
let ret = unsafe { self.next_unchecked() };
|
||||||
|
if ret.is_some() {
|
||||||
|
self.len -= 1
|
||||||
}
|
}
|
||||||
|
ret
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
fn size_hint(&self) -> (usize, Option<usize>) {
|
fn size_hint(&self) -> (usize, Option<usize>) {
|
||||||
let len = self.dict.len().unwrap_or_default();
|
let len = self.len as usize;
|
||||||
(
|
(len, Some(len))
|
||||||
len.saturating_sub(self.pos as usize),
|
|
||||||
Some(len.saturating_sub(self.pos as usize)),
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -269,7 +297,34 @@ impl<'a> std::iter::IntoIterator for &'a PyDict {
|
||||||
type IntoIter = PyDictIterator<'a>;
|
type IntoIter = PyDictIterator<'a>;
|
||||||
|
|
||||||
fn into_iter(self) -> Self::IntoIter {
|
fn into_iter(self) -> Self::IntoIter {
|
||||||
self.iter()
|
PyDictIterator {
|
||||||
|
dict: self,
|
||||||
|
ppos: 0,
|
||||||
|
di_used: self._len(),
|
||||||
|
len: self._len(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'py> PyDictIterator<'py> {
|
||||||
|
/// Advances the iterator without checking for concurrent modification.
|
||||||
|
///
|
||||||
|
/// See [`PyDict_Next`](https://docs.python.org/3/c-api/dict.html#c.PyDict_Next)
|
||||||
|
/// for more information.
|
||||||
|
unsafe fn next_unchecked(&mut self) -> Option<(&'py PyAny, &'py PyAny)> {
|
||||||
|
let mut key: *mut ffi::PyObject = std::ptr::null_mut();
|
||||||
|
let mut value: *mut ffi::PyObject = std::ptr::null_mut();
|
||||||
|
|
||||||
|
if ffi::PyDict_Next(self.dict.as_ptr(), &mut self.ppos, &mut key, &mut value) != 0 {
|
||||||
|
let py = self.dict.py();
|
||||||
|
// PyDict_Next returns borrowed values; for safety must make them owned (see #890)
|
||||||
|
Some((
|
||||||
|
py.from_owned_ptr(ffi::_Py_NewRef(key)),
|
||||||
|
py.from_owned_ptr(ffi::_Py_NewRef(value)),
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -660,6 +715,74 @@ mod tests {
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_iter_value_mutated() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let mut v = HashMap::new();
|
||||||
|
v.insert(7, 32);
|
||||||
|
v.insert(8, 42);
|
||||||
|
v.insert(9, 123);
|
||||||
|
|
||||||
|
let ob = v.to_object(py);
|
||||||
|
let dict = <PyDict as PyTryFrom>::try_from(ob.as_ref(py)).unwrap();
|
||||||
|
|
||||||
|
for (key, value) in dict.iter() {
|
||||||
|
dict.set_item(key, value.extract::<i32>().unwrap() + 7)
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[should_panic]
|
||||||
|
fn test_iter_key_mutated() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let mut v = HashMap::new();
|
||||||
|
for i in 0..10 {
|
||||||
|
v.insert(i * 2, i * 2);
|
||||||
|
}
|
||||||
|
let ob = v.to_object(py);
|
||||||
|
let dict = <PyDict as PyTryFrom>::try_from(ob.as_ref(py)).unwrap();
|
||||||
|
|
||||||
|
for (i, (key, value)) in dict.iter().enumerate() {
|
||||||
|
let key = key.extract::<i32>().unwrap();
|
||||||
|
let value = value.extract::<i32>().unwrap();
|
||||||
|
|
||||||
|
dict.set_item(key + 1, value + 1).unwrap();
|
||||||
|
|
||||||
|
if i > 1000 {
|
||||||
|
// avoid this test just running out of memory if it fails
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[should_panic]
|
||||||
|
fn test_iter_key_mutated_constant_len() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let mut v = HashMap::new();
|
||||||
|
for i in 0..10 {
|
||||||
|
v.insert(i * 2, i * 2);
|
||||||
|
}
|
||||||
|
let ob = v.to_object(py);
|
||||||
|
let dict = <PyDict as PyTryFrom>::try_from(ob.as_ref(py)).unwrap();
|
||||||
|
|
||||||
|
for (i, (key, value)) in dict.iter().enumerate() {
|
||||||
|
let key = key.extract::<i32>().unwrap();
|
||||||
|
let value = value.extract::<i32>().unwrap();
|
||||||
|
dict.del_item(key).unwrap();
|
||||||
|
dict.set_item(key + 1, value + 1).unwrap();
|
||||||
|
|
||||||
|
if i > 1000 {
|
||||||
|
// avoid this test just running out of memory if it fails
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_iter_size_hint() {
|
fn test_iter_size_hint() {
|
||||||
Python::with_gil(|py| {
|
Python::with_gil(|py| {
|
||||||
|
@ -675,10 +798,14 @@ mod tests {
|
||||||
iter.next();
|
iter.next();
|
||||||
assert_eq!(iter.size_hint(), (v.len() - 1, Some(v.len() - 1)));
|
assert_eq!(iter.size_hint(), (v.len() - 1, Some(v.len() - 1)));
|
||||||
|
|
||||||
// Exhust iterator.
|
// Exhaust iterator.
|
||||||
for _ in &mut iter {}
|
for _ in &mut iter {}
|
||||||
|
|
||||||
assert_eq!(iter.size_hint(), (0, Some(0)));
|
assert_eq!(iter.size_hint(), (0, Some(0)));
|
||||||
|
|
||||||
|
assert!(iter.next().is_none());
|
||||||
|
|
||||||
|
assert_eq!(iter.size_hint(), (0, Some(0)));
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -0,0 +1,194 @@
|
||||||
|
// Copyright (c) 2017-present PyO3 Project and Contributors
|
||||||
|
//
|
||||||
|
|
||||||
|
use crate::err::{PyErr, PyResult};
|
||||||
|
#[cfg(Py_LIMITED_API)]
|
||||||
|
use crate::types::PyIterator;
|
||||||
|
use crate::{ffi, AsPyPointer, PyAny, Python, ToPyObject};
|
||||||
|
|
||||||
|
use std::ptr;
|
||||||
|
|
||||||
|
/// Represents a Python `frozenset`
|
||||||
|
#[repr(transparent)]
|
||||||
|
pub struct PyFrozenSet(PyAny);
|
||||||
|
|
||||||
|
pyobject_native_type!(
|
||||||
|
PyFrozenSet,
|
||||||
|
ffi::PySetObject,
|
||||||
|
ffi::PyFrozenSet_Type,
|
||||||
|
#checkfunction=ffi::PyFrozenSet_Check
|
||||||
|
);
|
||||||
|
|
||||||
|
impl PyFrozenSet {
|
||||||
|
/// Creates a new frozenset.
|
||||||
|
///
|
||||||
|
/// May panic when running out of memory.
|
||||||
|
pub fn new<'p, T: ToPyObject>(py: Python<'p>, elements: &[T]) -> PyResult<&'p PyFrozenSet> {
|
||||||
|
let list = elements.to_object(py);
|
||||||
|
unsafe { py.from_owned_ptr_or_err(ffi::PyFrozenSet_New(list.as_ptr())) }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Creates a new empty frozen set
|
||||||
|
pub fn empty(py: Python<'_>) -> PyResult<&PyFrozenSet> {
|
||||||
|
unsafe { py.from_owned_ptr_or_err(ffi::PyFrozenSet_New(ptr::null_mut())) }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Return the number of items in the set.
|
||||||
|
/// This is equivalent to len(p) on a set.
|
||||||
|
#[inline]
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
unsafe { ffi::PySet_Size(self.as_ptr()) as usize }
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if set is empty.
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.len() == 0
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Determine if the set contains the specified key.
|
||||||
|
/// This is equivalent to the Python expression `key in self`.
|
||||||
|
pub fn contains<K>(&self, key: K) -> PyResult<bool>
|
||||||
|
where
|
||||||
|
K: ToPyObject,
|
||||||
|
{
|
||||||
|
unsafe {
|
||||||
|
match ffi::PySet_Contains(self.as_ptr(), key.to_object(self.py()).as_ptr()) {
|
||||||
|
1 => Ok(true),
|
||||||
|
0 => Ok(false),
|
||||||
|
_ => Err(PyErr::fetch(self.py())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns an iterator of values in this frozen set.
|
||||||
|
pub fn iter(&self) -> PyFrozenSetIterator<'_> {
|
||||||
|
IntoIterator::into_iter(self)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(Py_LIMITED_API)]
|
||||||
|
mod impl_ {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
|
||||||
|
type Item = &'a PyAny;
|
||||||
|
type IntoIter = PyFrozenSetIterator<'a>;
|
||||||
|
|
||||||
|
fn into_iter(self) -> Self::IntoIter {
|
||||||
|
PyFrozenSetIterator {
|
||||||
|
it: PyIterator::from_object(self.py(), self).unwrap(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct PyFrozenSetIterator<'p> {
|
||||||
|
it: &'p PyIterator,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'py> Iterator for PyFrozenSetIterator<'py> {
|
||||||
|
type Item = &'py super::PyAny;
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
|
self.it.next().map(Result::unwrap)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(not(Py_LIMITED_API))]
|
||||||
|
mod impl_ {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
|
||||||
|
type Item = &'a PyAny;
|
||||||
|
type IntoIter = PyFrozenSetIterator<'a>;
|
||||||
|
|
||||||
|
fn into_iter(self) -> Self::IntoIter {
|
||||||
|
PyFrozenSetIterator { set: self, pos: 0 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct PyFrozenSetIterator<'py> {
|
||||||
|
set: &'py PyAny,
|
||||||
|
pos: ffi::Py_ssize_t,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'py> Iterator for PyFrozenSetIterator<'py> {
|
||||||
|
type Item = &'py PyAny;
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
|
unsafe {
|
||||||
|
let mut key: *mut ffi::PyObject = std::ptr::null_mut();
|
||||||
|
let mut hash: ffi::Py_hash_t = 0;
|
||||||
|
if ffi::_PySet_NextEntry(self.set.as_ptr(), &mut self.pos, &mut key, &mut hash) != 0
|
||||||
|
{
|
||||||
|
// _PySet_NextEntry returns borrowed object; for safety must make owned (see #890)
|
||||||
|
Some(self.set.py().from_owned_ptr(ffi::_Py_NewRef(key)))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn size_hint(&self) -> (usize, Option<usize>) {
|
||||||
|
let len = self.set.len().unwrap_or_default();
|
||||||
|
(
|
||||||
|
len.saturating_sub(self.pos as usize),
|
||||||
|
Some(len.saturating_sub(self.pos as usize)),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub use impl_::*;
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_frozenset_new_and_len() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let set = PyFrozenSet::new(py, &[1]).unwrap();
|
||||||
|
assert_eq!(1, set.len());
|
||||||
|
|
||||||
|
let v = vec![1];
|
||||||
|
assert!(PyFrozenSet::new(py, &[v]).is_err());
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_frozenset_empty() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let set = PyFrozenSet::empty(py).unwrap();
|
||||||
|
assert_eq!(0, set.len());
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_frozenset_contains() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let set = PyFrozenSet::new(py, &[1]).unwrap();
|
||||||
|
assert!(set.contains(1).unwrap());
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_frozenset_iter() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let set = PyFrozenSet::new(py, &[1]).unwrap();
|
||||||
|
|
||||||
|
// iter method
|
||||||
|
for el in set.iter() {
|
||||||
|
assert_eq!(1i32, el.extract::<i32>().unwrap());
|
||||||
|
}
|
||||||
|
|
||||||
|
// intoiterator iteration
|
||||||
|
for el in set {
|
||||||
|
assert_eq!(1i32, el.extract::<i32>().unwrap());
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
|
@ -15,6 +15,7 @@ pub use self::datetime::{
|
||||||
};
|
};
|
||||||
pub use self::dict::{IntoPyDict, PyDict};
|
pub use self::dict::{IntoPyDict, PyDict};
|
||||||
pub use self::floatob::PyFloat;
|
pub use self::floatob::PyFloat;
|
||||||
|
pub use self::frozenset::PyFrozenSet;
|
||||||
pub use self::function::{PyCFunction, PyFunction};
|
pub use self::function::{PyCFunction, PyFunction};
|
||||||
pub use self::iterator::PyIterator;
|
pub use self::iterator::PyIterator;
|
||||||
pub use self::list::PyList;
|
pub use self::list::PyList;
|
||||||
|
@ -23,7 +24,7 @@ pub use self::module::PyModule;
|
||||||
pub use self::num::PyLong;
|
pub use self::num::PyLong;
|
||||||
pub use self::num::PyLong as PyInt;
|
pub use self::num::PyLong as PyInt;
|
||||||
pub use self::sequence::PySequence;
|
pub use self::sequence::PySequence;
|
||||||
pub use self::set::{PyFrozenSet, PySet};
|
pub use self::set::PySet;
|
||||||
pub use self::slice::{PySlice, PySliceIndices};
|
pub use self::slice::{PySlice, PySliceIndices};
|
||||||
#[cfg(all(not(Py_LIMITED_API), target_endian = "little"))]
|
#[cfg(all(not(Py_LIMITED_API), target_endian = "little"))]
|
||||||
pub use self::string::PyStringData;
|
pub use self::string::PyStringData;
|
||||||
|
@ -32,6 +33,42 @@ pub use self::traceback::PyTraceback;
|
||||||
pub use self::tuple::PyTuple;
|
pub use self::tuple::PyTuple;
|
||||||
pub use self::typeobject::PyType;
|
pub use self::typeobject::PyType;
|
||||||
|
|
||||||
|
/// Iteration over Python collections.
|
||||||
|
///
|
||||||
|
/// When working with a Python collection, one approach is to convert it to a Rust collection such
|
||||||
|
/// as `Vec` or `HashMap`. However this is a relatively expensive operation. If you just want to
|
||||||
|
/// visit all their items, consider iterating over the collections directly:
|
||||||
|
///
|
||||||
|
/// # Examples
|
||||||
|
///
|
||||||
|
/// ```rust
|
||||||
|
/// use pyo3::prelude::*;
|
||||||
|
/// use pyo3::types::PyDict;
|
||||||
|
///
|
||||||
|
/// # pub fn main() -> PyResult<()> {
|
||||||
|
/// Python::with_gil(|py| {
|
||||||
|
/// let dict: &PyDict = py.eval("{'a':'b', 'c':'d'}", None, None)?.cast_as()?;
|
||||||
|
///
|
||||||
|
/// for (key, value) in dict {
|
||||||
|
/// println!("key: {}, value: {}", key, value);
|
||||||
|
/// }
|
||||||
|
///
|
||||||
|
/// Ok(())
|
||||||
|
/// })
|
||||||
|
/// # }
|
||||||
|
/// ```
|
||||||
|
///
|
||||||
|
/// If PyO3 detects that the collection is mutated during iteration, it will panic.
|
||||||
|
///
|
||||||
|
/// These iterators use Python's C-API directly. However in certain cases, like when compiling for
|
||||||
|
/// the Limited API and PyPy, the underlying structures are opaque and that may not be possible.
|
||||||
|
/// In these cases the iterators are implemented by forwarding to [`PyIterator`].
|
||||||
|
pub mod iter {
|
||||||
|
pub use super::dict::PyDictIterator;
|
||||||
|
pub use super::frozenset::PyFrozenSetIterator;
|
||||||
|
pub use super::set::PySetIterator;
|
||||||
|
}
|
||||||
|
|
||||||
// Implementations core to all native types
|
// Implementations core to all native types
|
||||||
#[doc(hidden)]
|
#[doc(hidden)]
|
||||||
#[macro_export]
|
#[macro_export]
|
||||||
|
@ -225,6 +262,7 @@ mod complex;
|
||||||
mod datetime;
|
mod datetime;
|
||||||
mod dict;
|
mod dict;
|
||||||
mod floatob;
|
mod floatob;
|
||||||
|
mod frozenset;
|
||||||
mod function;
|
mod function;
|
||||||
mod iterator;
|
mod iterator;
|
||||||
mod list;
|
mod list;
|
||||||
|
|
296
src/types/set.rs
296
src/types/set.rs
|
@ -13,17 +13,7 @@ use std::{collections, hash, ptr};
|
||||||
#[repr(transparent)]
|
#[repr(transparent)]
|
||||||
pub struct PySet(PyAny);
|
pub struct PySet(PyAny);
|
||||||
|
|
||||||
/// Represents a Python `frozenset`
|
|
||||||
#[repr(transparent)]
|
|
||||||
pub struct PyFrozenSet(PyAny);
|
|
||||||
|
|
||||||
pyobject_native_type!(PySet, ffi::PySetObject, ffi::PySet_Type, #checkfunction=ffi::PySet_Check);
|
pyobject_native_type!(PySet, ffi::PySetObject, ffi::PySet_Type, #checkfunction=ffi::PySet_Check);
|
||||||
pyobject_native_type!(
|
|
||||||
PyFrozenSet,
|
|
||||||
ffi::PySetObject,
|
|
||||||
ffi::PyFrozenSet_Type,
|
|
||||||
#checkfunction=ffi::PyFrozenSet_Check
|
|
||||||
);
|
|
||||||
|
|
||||||
impl PySet {
|
impl PySet {
|
||||||
/// Creates a new set with elements from the given slice.
|
/// Creates a new set with elements from the given slice.
|
||||||
|
@ -111,86 +101,118 @@ impl PySet {
|
||||||
|
|
||||||
/// Returns an iterator of values in this set.
|
/// Returns an iterator of values in this set.
|
||||||
///
|
///
|
||||||
/// Note that it can be unsafe to use when the set might be changed by other code.
|
/// # Panics
|
||||||
|
///
|
||||||
|
/// If PyO3 detects that the set is mutated during iteration, it will panic.
|
||||||
pub fn iter(&self) -> PySetIterator<'_> {
|
pub fn iter(&self) -> PySetIterator<'_> {
|
||||||
PySetIterator::new(self)
|
IntoIterator::into_iter(self)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(Py_LIMITED_API)]
|
#[cfg(Py_LIMITED_API)]
|
||||||
pub struct PySetIterator<'p> {
|
mod impl_ {
|
||||||
it: &'p PyIterator,
|
use super::*;
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(Py_LIMITED_API)]
|
impl<'a> std::iter::IntoIterator for &'a PySet {
|
||||||
impl PySetIterator<'_> {
|
type Item = &'a PyAny;
|
||||||
fn new(set: &PyAny) -> PySetIterator<'_> {
|
type IntoIter = PySetIterator<'a>;
|
||||||
PySetIterator {
|
|
||||||
it: PyIterator::from_object(set.py(), set).unwrap(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(Py_LIMITED_API)]
|
/// Returns an iterator of values in this set.
|
||||||
impl<'py> Iterator for PySetIterator<'py> {
|
///
|
||||||
type Item = &'py super::PyAny;
|
/// # Panics
|
||||||
|
///
|
||||||
#[inline]
|
/// If PyO3 detects that the set is mutated during iteration, it will panic.
|
||||||
fn next(&mut self) -> Option<Self::Item> {
|
fn into_iter(self) -> Self::IntoIter {
|
||||||
self.it.next().map(|p| p.unwrap())
|
PySetIterator {
|
||||||
}
|
it: PyIterator::from_object(self.py(), self).unwrap(),
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(not(Py_LIMITED_API))]
|
|
||||||
pub struct PySetIterator<'py> {
|
|
||||||
set: &'py super::PyAny,
|
|
||||||
pos: isize,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(not(Py_LIMITED_API))]
|
|
||||||
impl PySetIterator<'_> {
|
|
||||||
fn new(set: &PyAny) -> PySetIterator<'_> {
|
|
||||||
PySetIterator { set, pos: 0 }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(not(Py_LIMITED_API))]
|
|
||||||
impl<'py> Iterator for PySetIterator<'py> {
|
|
||||||
type Item = &'py super::PyAny;
|
|
||||||
|
|
||||||
#[inline]
|
|
||||||
fn next(&mut self) -> Option<Self::Item> {
|
|
||||||
unsafe {
|
|
||||||
let mut key: *mut ffi::PyObject = std::ptr::null_mut();
|
|
||||||
let mut hash: ffi::Py_hash_t = 0;
|
|
||||||
if ffi::_PySet_NextEntry(self.set.as_ptr(), &mut self.pos, &mut key, &mut hash) != 0 {
|
|
||||||
// _PySet_NextEntry returns borrowed object; for safety must make owned (see #890)
|
|
||||||
Some(self.set.py().from_owned_ptr(ffi::_Py_NewRef(key)))
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
pub struct PySetIterator<'p> {
|
||||||
fn size_hint(&self) -> (usize, Option<usize>) {
|
it: &'p PyIterator,
|
||||||
let len = self.set.len().unwrap_or_default();
|
}
|
||||||
(
|
|
||||||
len.saturating_sub(self.pos as usize),
|
impl<'py> Iterator for PySetIterator<'py> {
|
||||||
Some(len.saturating_sub(self.pos as usize)),
|
type Item = &'py super::PyAny;
|
||||||
)
|
|
||||||
|
/// Advances the iterator and returns the next value.
|
||||||
|
///
|
||||||
|
/// # Panics
|
||||||
|
///
|
||||||
|
/// If PyO3 detects that the set is mutated during iteration, it will panic.
|
||||||
|
#[inline]
|
||||||
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
|
self.it.next().map(Result::unwrap)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<'a> std::iter::IntoIterator for &'a PySet {
|
#[cfg(not(Py_LIMITED_API))]
|
||||||
type Item = &'a PyAny;
|
mod impl_ {
|
||||||
type IntoIter = PySetIterator<'a>;
|
use super::*;
|
||||||
|
pub struct PySetIterator<'py> {
|
||||||
|
set: &'py super::PyAny,
|
||||||
|
pos: ffi::Py_ssize_t,
|
||||||
|
used: ffi::Py_ssize_t,
|
||||||
|
}
|
||||||
|
|
||||||
fn into_iter(self) -> Self::IntoIter {
|
impl<'a> std::iter::IntoIterator for &'a PySet {
|
||||||
self.iter()
|
type Item = &'a PyAny;
|
||||||
|
type IntoIter = PySetIterator<'a>;
|
||||||
|
/// Returns an iterator of values in this set.
|
||||||
|
///
|
||||||
|
/// # Panics
|
||||||
|
///
|
||||||
|
/// If PyO3 detects that the set is mutated during iteration, it will panic.
|
||||||
|
fn into_iter(self) -> Self::IntoIter {
|
||||||
|
PySetIterator {
|
||||||
|
set: self,
|
||||||
|
pos: 0,
|
||||||
|
used: unsafe { ffi::PySet_Size(self.as_ptr()) },
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'py> Iterator for PySetIterator<'py> {
|
||||||
|
type Item = &'py super::PyAny;
|
||||||
|
|
||||||
|
/// Advances the iterator and returns the next value.
|
||||||
|
///
|
||||||
|
/// # Panics
|
||||||
|
///
|
||||||
|
/// If PyO3 detects that the set is mutated during iteration, it will panic.
|
||||||
|
#[inline]
|
||||||
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
|
unsafe {
|
||||||
|
let len = ffi::PySet_Size(self.set.as_ptr());
|
||||||
|
assert_eq!(self.used, len, "Set changed size during iteration");
|
||||||
|
|
||||||
|
let mut key: *mut ffi::PyObject = std::ptr::null_mut();
|
||||||
|
let mut hash: ffi::Py_hash_t = 0;
|
||||||
|
if ffi::_PySet_NextEntry(self.set.as_ptr(), &mut self.pos, &mut key, &mut hash) != 0
|
||||||
|
{
|
||||||
|
// _PySet_NextEntry returns borrowed object; for safety must make owned (see #890)
|
||||||
|
Some(self.set.py().from_owned_ptr(ffi::_Py_NewRef(key)))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn size_hint(&self) -> (usize, Option<usize>) {
|
||||||
|
let len = self.set.len().unwrap_or_default();
|
||||||
|
(
|
||||||
|
len.saturating_sub(self.pos as usize),
|
||||||
|
Some(len.saturating_sub(self.pos as usize)),
|
||||||
|
)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub use impl_::*;
|
||||||
|
|
||||||
impl<T, S> ToPyObject for collections::HashSet<T, S>
|
impl<T, S> ToPyObject for collections::HashSet<T, S>
|
||||||
where
|
where
|
||||||
T: hash::Hash + Eq + ToPyObject,
|
T: hash::Hash + Eq + ToPyObject,
|
||||||
|
@ -274,67 +296,9 @@ where
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl PyFrozenSet {
|
|
||||||
/// Creates a new frozenset.
|
|
||||||
///
|
|
||||||
/// May panic when running out of memory.
|
|
||||||
pub fn new<'p, T: ToPyObject>(py: Python<'p>, elements: &[T]) -> PyResult<&'p PyFrozenSet> {
|
|
||||||
let list = elements.to_object(py);
|
|
||||||
unsafe { py.from_owned_ptr_or_err(ffi::PyFrozenSet_New(list.as_ptr())) }
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Creates a new empty frozen set
|
|
||||||
pub fn empty(py: Python<'_>) -> PyResult<&PyFrozenSet> {
|
|
||||||
unsafe { py.from_owned_ptr_or_err(ffi::PyFrozenSet_New(ptr::null_mut())) }
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Return the number of items in the set.
|
|
||||||
/// This is equivalent to len(p) on a set.
|
|
||||||
#[inline]
|
|
||||||
pub fn len(&self) -> usize {
|
|
||||||
unsafe { ffi::PySet_Size(self.as_ptr()) as usize }
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Check if set is empty.
|
|
||||||
pub fn is_empty(&self) -> bool {
|
|
||||||
self.len() == 0
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Determine if the set contains the specified key.
|
|
||||||
/// This is equivalent to the Python expression `key in self`.
|
|
||||||
pub fn contains<K>(&self, key: K) -> PyResult<bool>
|
|
||||||
where
|
|
||||||
K: ToPyObject,
|
|
||||||
{
|
|
||||||
unsafe {
|
|
||||||
match ffi::PySet_Contains(self.as_ptr(), key.to_object(self.py()).as_ptr()) {
|
|
||||||
1 => Ok(true),
|
|
||||||
0 => Ok(false),
|
|
||||||
_ => Err(PyErr::fetch(self.py())),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Returns an iterator of values in this frozen set.
|
|
||||||
///
|
|
||||||
/// Note that it can be unsafe to use when the set might be changed by other code.
|
|
||||||
pub fn iter(&self) -> PySetIterator<'_> {
|
|
||||||
PySetIterator::new(self.as_ref())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
|
|
||||||
type Item = &'a PyAny;
|
|
||||||
type IntoIter = PySetIterator<'a>;
|
|
||||||
|
|
||||||
fn into_iter(self) -> Self::IntoIter {
|
|
||||||
self.iter()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{PyFrozenSet, PySet};
|
use super::PySet;
|
||||||
use crate::{IntoPy, PyObject, PyTryFrom, Python, ToPyObject};
|
use crate::{IntoPy, PyObject, PyTryFrom, Python, ToPyObject};
|
||||||
use std::collections::{BTreeSet, HashSet};
|
use std::collections::{BTreeSet, HashSet};
|
||||||
|
|
||||||
|
@ -440,6 +404,32 @@ mod tests {
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[should_panic]
|
||||||
|
fn test_set_iter_mutation() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let set = PySet::new(py, &[1, 2, 3, 4, 5]).unwrap();
|
||||||
|
|
||||||
|
for _ in set {
|
||||||
|
let _ = set.add(42);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[should_panic]
|
||||||
|
fn test_set_iter_mutation_same_len() {
|
||||||
|
Python::with_gil(|py| {
|
||||||
|
let set = PySet::new(py, &[1, 2, 3, 4, 5]).unwrap();
|
||||||
|
|
||||||
|
for item in set {
|
||||||
|
let item: i32 = item.extract().unwrap();
|
||||||
|
let _ = set.del_item(item);
|
||||||
|
let _ = set.add(item + 10);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_set_iter_size_hint() {
|
fn test_set_iter_size_hint() {
|
||||||
Python::with_gil(|py| {
|
Python::with_gil(|py| {
|
||||||
|
@ -457,50 +447,6 @@ mod tests {
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_frozenset_new_and_len() {
|
|
||||||
Python::with_gil(|py| {
|
|
||||||
let set = PyFrozenSet::new(py, &[1]).unwrap();
|
|
||||||
assert_eq!(1, set.len());
|
|
||||||
|
|
||||||
let v = vec![1];
|
|
||||||
assert!(PyFrozenSet::new(py, &[v]).is_err());
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_frozenset_empty() {
|
|
||||||
Python::with_gil(|py| {
|
|
||||||
let set = PyFrozenSet::empty(py).unwrap();
|
|
||||||
assert_eq!(0, set.len());
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_frozenset_contains() {
|
|
||||||
Python::with_gil(|py| {
|
|
||||||
let set = PyFrozenSet::new(py, &[1]).unwrap();
|
|
||||||
assert!(set.contains(1).unwrap());
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_frozenset_iter() {
|
|
||||||
Python::with_gil(|py| {
|
|
||||||
let set = PyFrozenSet::new(py, &[1]).unwrap();
|
|
||||||
|
|
||||||
// iter method
|
|
||||||
for el in set.iter() {
|
|
||||||
assert_eq!(1i32, el.extract::<i32>().unwrap());
|
|
||||||
}
|
|
||||||
|
|
||||||
// intoiterator iteration
|
|
||||||
for el in set {
|
|
||||||
assert_eq!(1i32, el.extract::<i32>().unwrap());
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_extract_hashset() {
|
fn test_extract_hashset() {
|
||||||
Python::with_gil(|py| {
|
Python::with_gil(|py| {
|
||||||
|
|
Loading…
Reference in New Issue