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