Skip to main content

pyo3/conversions/std/
set.rs

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