pyo3/conversions/std/
set.rs

1use std::{cmp, collections, hash};
2
3#[cfg(feature = "experimental-inspect")]
4use crate::inspect::types::TypeInfo;
5use crate::{
6    instance::Bound,
7    types::any::PyAnyMethods,
8    types::frozenset::PyFrozenSetMethods,
9    types::set::{new_from_iter, PySetMethods},
10    types::{PyFrozenSet, PySet},
11    FromPyObject, IntoPy, PyAny, PyErr, PyObject, PyResult, Python, ToPyObject,
12};
13
14impl<T, S> ToPyObject for collections::HashSet<T, S>
15where
16    T: hash::Hash + Eq + ToPyObject,
17    S: hash::BuildHasher + Default,
18{
19    fn to_object(&self, py: Python<'_>) -> PyObject {
20        new_from_iter(py, self)
21            .expect("Failed to create Python set from HashSet")
22            .into()
23    }
24}
25
26impl<T> ToPyObject for collections::BTreeSet<T>
27where
28    T: hash::Hash + Eq + ToPyObject,
29{
30    fn to_object(&self, py: Python<'_>) -> PyObject {
31        new_from_iter(py, self)
32            .expect("Failed to create Python set from BTreeSet")
33            .into()
34    }
35}
36
37impl<K, S> IntoPy<PyObject> for collections::HashSet<K, S>
38where
39    K: IntoPy<PyObject> + Eq + hash::Hash,
40    S: hash::BuildHasher + Default,
41{
42    fn into_py(self, py: Python<'_>) -> PyObject {
43        new_from_iter(py, self.into_iter().map(|item| item.into_py(py)))
44            .expect("Failed to create Python set from HashSet")
45            .into()
46    }
47
48    #[cfg(feature = "experimental-inspect")]
49    fn type_output() -> TypeInfo {
50        TypeInfo::set_of(K::type_output())
51    }
52}
53
54impl<'py, K, S> FromPyObject<'py> for collections::HashSet<K, S>
55where
56    K: FromPyObject<'py> + cmp::Eq + hash::Hash,
57    S: hash::BuildHasher + Default,
58{
59    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
60        match ob.downcast::<PySet>() {
61            Ok(set) => set.iter().map(|any| any.extract()).collect(),
62            Err(err) => {
63                if let Ok(frozen_set) = ob.downcast::<PyFrozenSet>() {
64                    frozen_set.iter().map(|any| any.extract()).collect()
65                } else {
66                    Err(PyErr::from(err))
67                }
68            }
69        }
70    }
71
72    #[cfg(feature = "experimental-inspect")]
73    fn type_input() -> TypeInfo {
74        TypeInfo::set_of(K::type_input())
75    }
76}
77
78impl<K> IntoPy<PyObject> for collections::BTreeSet<K>
79where
80    K: IntoPy<PyObject> + cmp::Ord,
81{
82    fn into_py(self, py: Python<'_>) -> PyObject {
83        new_from_iter(py, self.into_iter().map(|item| item.into_py(py)))
84            .expect("Failed to create Python set from BTreeSet")
85            .into()
86    }
87
88    #[cfg(feature = "experimental-inspect")]
89    fn type_output() -> TypeInfo {
90        TypeInfo::set_of(K::type_output())
91    }
92}
93
94impl<'py, K> FromPyObject<'py> for collections::BTreeSet<K>
95where
96    K: FromPyObject<'py> + cmp::Ord,
97{
98    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
99        match ob.downcast::<PySet>() {
100            Ok(set) => set.iter().map(|any| any.extract()).collect(),
101            Err(err) => {
102                if let Ok(frozen_set) = ob.downcast::<PyFrozenSet>() {
103                    frozen_set.iter().map(|any| any.extract()).collect()
104                } else {
105                    Err(PyErr::from(err))
106                }
107            }
108        }
109    }
110
111    #[cfg(feature = "experimental-inspect")]
112    fn type_input() -> TypeInfo {
113        TypeInfo::set_of(K::type_input())
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use crate::types::{any::PyAnyMethods, PyFrozenSet, PySet};
120    use crate::{IntoPy, PyObject, Python, ToPyObject};
121    use std::collections::{BTreeSet, HashSet};
122
123    #[test]
124    fn test_extract_hashset() {
125        Python::with_gil(|py| {
126            let set = PySet::new_bound(py, &[1, 2, 3, 4, 5]).unwrap();
127            let hash_set: HashSet<usize> = set.extract().unwrap();
128            assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
129
130            let set = PyFrozenSet::new_bound(py, &[1, 2, 3, 4, 5]).unwrap();
131            let hash_set: HashSet<usize> = set.extract().unwrap();
132            assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
133        });
134    }
135
136    #[test]
137    fn test_extract_btreeset() {
138        Python::with_gil(|py| {
139            let set = PySet::new_bound(py, &[1, 2, 3, 4, 5]).unwrap();
140            let hash_set: BTreeSet<usize> = set.extract().unwrap();
141            assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
142
143            let set = PyFrozenSet::new_bound(py, &[1, 2, 3, 4, 5]).unwrap();
144            let hash_set: BTreeSet<usize> = set.extract().unwrap();
145            assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
146        });
147    }
148
149    #[test]
150    fn test_set_into_py() {
151        Python::with_gil(|py| {
152            let bt: BTreeSet<u64> = [1, 2, 3, 4, 5].iter().cloned().collect();
153            let hs: HashSet<u64> = [1, 2, 3, 4, 5].iter().cloned().collect();
154
155            let bto: PyObject = bt.clone().into_py(py);
156            let hso: PyObject = hs.clone().into_py(py);
157
158            assert_eq!(bt, bto.extract(py).unwrap());
159            assert_eq!(hs, hso.extract(py).unwrap());
160        });
161    }
162
163    #[test]
164    fn test_set_to_object() {
165        Python::with_gil(|py| {
166            let bt: BTreeSet<u64> = [1, 2, 3, 4, 5].iter().cloned().collect();
167            let hs: HashSet<u64> = [1, 2, 3, 4, 5].iter().cloned().collect();
168
169            let bto: PyObject = bt.to_object(py);
170            let hso: PyObject = hs.to_object(py);
171
172            assert_eq!(bt, bto.extract(py).unwrap());
173            assert_eq!(hs, hso.extract(py).unwrap());
174        });
175    }
176}
pyo3/conversions/std/set.rs

pyo3/conversions/std/
set.rs
