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.
19use 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}