luxolus
/
pyo3
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #3632 from messense/extract-frozen-set 

Add support for extracting Rust set types from `frozenset`
This commit is contained in:
David Hewitt 2023-12-07 16:17:51 +00:00 committed by GitHub
parent 601d9573c8 c21a84d999
commit 24d9113974
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 48 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
newsfragments/3632.added.md Normal file
View File

@ -0,0 +1 @@
Add support for extracting Rust set types from `frozenset`.

18
src/conversions/hashbrown.rs
View File

@ -18,7 +18,7 @@
//! The required hashbrown version may vary based on the version of PyO3.
use crate::{
types::set::new_from_iter,
types::{IntoPyDict, PyDict, PySet},
types::{IntoPyDict, PyDict, PyFrozenSet, PySet},
FromPyObject, IntoPy, PyAny, PyErr, PyObject, PyResult, Python, ToPyObject,
};
use std::{cmp, hash};
@ -93,8 +93,16 @@ where
S: hash::BuildHasher + Default,
{
fn extract(ob: &'source PyAny) -> PyResult<Self> {
let set: &PySet = ob.downcast()?;
set.iter().map(K::extract).collect()
match ob.downcast::<PySet>() {
Ok(set) => set.iter().map(K::extract).collect(),
Err(err) => {
if let Ok(frozen_set) = ob.downcast::<PyFrozenSet>() {
frozen_set.iter().map(K::extract).collect()
} else {
Err(PyErr::from(err))
}
}
}
}
}
@ -173,6 +181,10 @@ mod tests {
let set = PySet::new(py, &[1, 2, 3, 4, 5]).unwrap();
let hash_set: hashbrown::HashSet<usize> = set.extract().unwrap();
assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
let set = PyFrozenSet::new(py, &[1, 2, 3, 4, 5]).unwrap();
let hash_set: hashbrown::HashSet<usize> = set.extract().unwrap();
assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
});
}

39
src/conversions/std/set.rs
View File

@ -3,8 +3,9 @@ use std::{cmp, collections, hash};
#[cfg(feature = "experimental-inspect")]
use crate::inspect::types::TypeInfo;
use crate::{
types::set::new_from_iter, types::PySet, FromPyObject, IntoPy, PyAny, PyObject, PyResult,
Python, ToPyObject,
types::set::new_from_iter,
types::{PyFrozenSet, PySet},
FromPyObject, IntoPy, PyAny, PyErr, PyObject, PyResult, Python, ToPyObject,
};
impl<T, S> ToPyObject for collections::HashSet<T, S>
@ -53,8 +54,16 @@ where
S: hash::BuildHasher + Default,
{
fn extract(ob: &'source PyAny) -> PyResult<Self> {
let set: &PySet = ob.downcast()?;
set.iter().map(K::extract).collect()
match ob.downcast::<PySet>() {
Ok(set) => set.iter().map(K::extract).collect(),
Err(err) => {
if let Ok(frozen_set) = ob.downcast::<PyFrozenSet>() {
frozen_set.iter().map(K::extract).collect()
} else {
Err(PyErr::from(err))
}
}
}
}
#[cfg(feature = "experimental-inspect")]
@ -84,8 +93,16 @@ where
K: FromPyObject<'source> + cmp::Ord,
{
fn extract(ob: &'source PyAny) -> PyResult<Self> {
let set: &PySet = ob.downcast()?;
set.iter().map(K::extract).collect()
match ob.downcast::<PySet>() {
Ok(set) => set.iter().map(K::extract).collect(),
Err(err) => {
if let Ok(frozen_set) = ob.downcast::<PyFrozenSet>() {
frozen_set.iter().map(K::extract).collect()
} else {
Err(PyErr::from(err))
}
}
}
}
#[cfg(feature = "experimental-inspect")]
@ -96,7 +113,7 @@ where
#[cfg(test)]
mod tests {
use super::PySet;
use super::{PyFrozenSet, PySet};
use crate::{IntoPy, PyObject, Python, ToPyObject};
use std::collections::{BTreeSet, HashSet};
@ -106,6 +123,10 @@ mod tests {
let set = PySet::new(py, &[1, 2, 3, 4, 5]).unwrap();
let hash_set: HashSet<usize> = set.extract().unwrap();
assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
let set = PyFrozenSet::new(py, &[1, 2, 3, 4, 5]).unwrap();
let hash_set: HashSet<usize> = set.extract().unwrap();
assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
});
}
@ -115,6 +136,10 @@ mod tests {
let set = PySet::new(py, &[1, 2, 3, 4, 5]).unwrap();
let hash_set: BTreeSet<usize> = set.extract().unwrap();
assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
let set = PyFrozenSet::new(py, &[1, 2, 3, 4, 5]).unwrap();
let hash_set: BTreeSet<usize> = set.extract().unwrap();
assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
});
}