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