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

Add support for extracting Rust set types from `frozenset`
PyO3 · Dec 7, 2023 · 24d9113 · 24d9113
2 parents 601d957 + c21a84d
commit 24d9113
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 10 deletions.
diff --git a/newsfragments/3632.added.md b/newsfragments/3632.added.md
@@ -0,0 +1 @@
+Add support for extracting Rust set types from `frozenset`.
diff --git a/src/conversions/hashbrown.rs b/src/conversions/hashbrown.rs
@@ -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());
         });
     }
 

diff --git a/src/conversions/std/set.rs b/src/conversions/std/set.rs
@@ -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());
         });
     }