pyo3/conversions/std/
map.rs

1use std::{cmp, collections, hash};
2
3#[cfg(feature = "experimental-inspect")]
4use crate::inspect::types::TypeInfo;
5use crate::{
6    instance::Bound,
7    types::dict::PyDictMethods,
8    types::{any::PyAnyMethods, IntoPyDict, PyDict},
9    FromPyObject, IntoPy, PyAny, PyErr, PyObject, Python, ToPyObject,
10};
11
12impl<K, V, H> ToPyObject for collections::HashMap<K, V, H>
13where
14    K: hash::Hash + cmp::Eq + ToPyObject,
15    V: ToPyObject,
16    H: hash::BuildHasher,
17{
18    fn to_object(&self, py: Python<'_>) -> PyObject {
19        IntoPyDict::into_py_dict_bound(self, py).into()
20    }
21}
22
23impl<K, V> ToPyObject for collections::BTreeMap<K, V>
24where
25    K: cmp::Eq + ToPyObject,
26    V: ToPyObject,
27{
28    fn to_object(&self, py: Python<'_>) -> PyObject {
29        IntoPyDict::into_py_dict_bound(self, py).into()
30    }
31}
32
33impl<K, V, H> IntoPy<PyObject> for collections::HashMap<K, V, H>
34where
35    K: hash::Hash + cmp::Eq + IntoPy<PyObject>,
36    V: IntoPy<PyObject>,
37    H: hash::BuildHasher,
38{
39    fn into_py(self, py: Python<'_>) -> PyObject {
40        let iter = self
41            .into_iter()
42            .map(|(k, v)| (k.into_py(py), v.into_py(py)));
43        IntoPyDict::into_py_dict_bound(iter, py).into()
44    }
45
46    #[cfg(feature = "experimental-inspect")]
47    fn type_output() -> TypeInfo {
48        TypeInfo::dict_of(K::type_output(), V::type_output())
49    }
50}
51
52impl<K, V> IntoPy<PyObject> for collections::BTreeMap<K, V>
53where
54    K: cmp::Eq + IntoPy<PyObject>,
55    V: IntoPy<PyObject>,
56{
57    fn into_py(self, py: Python<'_>) -> PyObject {
58        let iter = self
59            .into_iter()
60            .map(|(k, v)| (k.into_py(py), v.into_py(py)));
61        IntoPyDict::into_py_dict_bound(iter, py).into()
62    }
63
64    #[cfg(feature = "experimental-inspect")]
65    fn type_output() -> TypeInfo {
66        TypeInfo::dict_of(K::type_output(), V::type_output())
67    }
68}
69
70impl<'py, K, V, S> FromPyObject<'py> for collections::HashMap<K, V, S>
71where
72    K: FromPyObject<'py> + cmp::Eq + hash::Hash,
73    V: FromPyObject<'py>,
74    S: hash::BuildHasher + Default,
75{
76    fn extract_bound(ob: &Bound<'py, PyAny>) -> Result<Self, PyErr> {
77        let dict = ob.downcast::<PyDict>()?;
78        let mut ret = collections::HashMap::with_capacity_and_hasher(dict.len(), S::default());
79        for (k, v) in dict {
80            ret.insert(k.extract()?, v.extract()?);
81        }
82        Ok(ret)
83    }
84
85    #[cfg(feature = "experimental-inspect")]
86    fn type_input() -> TypeInfo {
87        TypeInfo::mapping_of(K::type_input(), V::type_input())
88    }
89}
90
91impl<'py, K, V> FromPyObject<'py> for collections::BTreeMap<K, V>
92where
93    K: FromPyObject<'py> + cmp::Ord,
94    V: FromPyObject<'py>,
95{
96    fn extract_bound(ob: &Bound<'py, PyAny>) -> Result<Self, PyErr> {
97        let dict = ob.downcast::<PyDict>()?;
98        let mut ret = collections::BTreeMap::new();
99        for (k, v) in dict {
100            ret.insert(k.extract()?, v.extract()?);
101        }
102        Ok(ret)
103    }
104
105    #[cfg(feature = "experimental-inspect")]
106    fn type_input() -> TypeInfo {
107        TypeInfo::mapping_of(K::type_input(), V::type_input())
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use std::collections::{BTreeMap, HashMap};
115
116    #[test]
117    fn test_hashmap_to_python() {
118        Python::with_gil(|py| {
119            let mut map = HashMap::<i32, i32>::new();
120            map.insert(1, 1);
121
122            let m = map.to_object(py);
123            let py_map = m.downcast_bound::<PyDict>(py).unwrap();
124
125            assert!(py_map.len() == 1);
126            assert!(
127                py_map
128                    .get_item(1)
129                    .unwrap()
130                    .unwrap()
131                    .extract::<i32>()
132                    .unwrap()
133                    == 1
134            );
135            assert_eq!(map, py_map.extract().unwrap());
136        });
137    }
138
139    #[test]
140    fn test_btreemap_to_python() {
141        Python::with_gil(|py| {
142            let mut map = BTreeMap::<i32, i32>::new();
143            map.insert(1, 1);
144
145            let m = map.to_object(py);
146            let py_map = m.downcast_bound::<PyDict>(py).unwrap();
147
148            assert!(py_map.len() == 1);
149            assert!(
150                py_map
151                    .get_item(1)
152                    .unwrap()
153                    .unwrap()
154                    .extract::<i32>()
155                    .unwrap()
156                    == 1
157            );
158            assert_eq!(map, py_map.extract().unwrap());
159        });
160    }
161
162    #[test]
163    fn test_hashmap_into_python() {
164        Python::with_gil(|py| {
165            let mut map = HashMap::<i32, i32>::new();
166            map.insert(1, 1);
167
168            let m: PyObject = map.into_py(py);
169            let py_map = m.downcast_bound::<PyDict>(py).unwrap();
170
171            assert!(py_map.len() == 1);
172            assert!(
173                py_map
174                    .get_item(1)
175                    .unwrap()
176                    .unwrap()
177                    .extract::<i32>()
178                    .unwrap()
179                    == 1
180            );
181        });
182    }
183
184    #[test]
185    fn test_btreemap_into_py() {
186        Python::with_gil(|py| {
187            let mut map = BTreeMap::<i32, i32>::new();
188            map.insert(1, 1);
189
190            let m: PyObject = map.into_py(py);
191            let py_map = m.downcast_bound::<PyDict>(py).unwrap();
192
193            assert!(py_map.len() == 1);
194            assert!(
195                py_map
196                    .get_item(1)
197                    .unwrap()
198                    .unwrap()
199                    .extract::<i32>()
200                    .unwrap()
201                    == 1
202            );
203        });
204    }
205}