Skip to main content

pyo3/conversions/
hashbrown.rs

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