pyo3/conversions/std/
set.rs

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