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}