Skip to main content

pyo3/conversions/std/
map.rs

1#[cfg(feature = "experimental-inspect")]
2use crate::inspect::{type_hint_subscript, PyStaticExpr};
3#[cfg(feature = "experimental-inspect")]
4use crate::type_object::PyTypeInfo;
5use crate::{
6    conversion::{FromPyObjectOwned, IntoPyObject},
7    instance::Bound,
8    types::{any::PyAnyMethods, dict::PyDictMethods, PyDict},
9    Borrowed, FromPyObject, PyAny, PyErr, Python,
10};
11use std::{cmp, collections, hash};
12
13impl<'py, K, V, H> IntoPyObject<'py> for collections::HashMap<K, V, H>
14where
15    K: IntoPyObject<'py> + cmp::Eq + hash::Hash,
16    V: IntoPyObject<'py>,
17    H: hash::BuildHasher,
18{
19    type Target = PyDict;
20    type Output = Bound<'py, Self::Target>;
21    type Error = PyErr;
22
23    #[cfg(feature = "experimental-inspect")]
24    const OUTPUT_TYPE: PyStaticExpr =
25        type_hint_subscript!(PyDict::TYPE_HINT, K::OUTPUT_TYPE, V::OUTPUT_TYPE);
26
27    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
28        let dict = PyDict::new(py);
29        for (k, v) in self {
30            dict.set_item(k, v)?;
31        }
32        Ok(dict)
33    }
34}
35
36impl<'a, 'py, K, V, H> IntoPyObject<'py> for &'a collections::HashMap<K, V, H>
37where
38    &'a K: IntoPyObject<'py> + cmp::Eq + hash::Hash,
39    &'a V: IntoPyObject<'py>,
40    H: hash::BuildHasher,
41{
42    type Target = PyDict;
43    type Output = Bound<'py, Self::Target>;
44    type Error = PyErr;
45
46    #[cfg(feature = "experimental-inspect")]
47    const OUTPUT_TYPE: PyStaticExpr =
48        type_hint_subscript!(PyDict::TYPE_HINT, <&K>::OUTPUT_TYPE, <&V>::OUTPUT_TYPE);
49
50    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
51        let dict = PyDict::new(py);
52        for (k, v) in self {
53            dict.set_item(k, v)?;
54        }
55        Ok(dict)
56    }
57}
58
59impl<'py, K, V> IntoPyObject<'py> for collections::BTreeMap<K, V>
60where
61    K: IntoPyObject<'py> + cmp::Eq,
62    V: IntoPyObject<'py>,
63{
64    type Target = PyDict;
65    type Output = Bound<'py, Self::Target>;
66    type Error = PyErr;
67
68    #[cfg(feature = "experimental-inspect")]
69    const OUTPUT_TYPE: PyStaticExpr =
70        type_hint_subscript!(PyDict::TYPE_HINT, K::OUTPUT_TYPE, V::OUTPUT_TYPE);
71
72    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
73        let dict = PyDict::new(py);
74        for (k, v) in self {
75            dict.set_item(k, v)?;
76        }
77        Ok(dict)
78    }
79}
80
81impl<'a, 'py, K, V> IntoPyObject<'py> for &'a collections::BTreeMap<K, V>
82where
83    &'a K: IntoPyObject<'py> + cmp::Eq,
84    &'a V: IntoPyObject<'py>,
85    K: 'a,
86    V: 'a,
87{
88    type Target = PyDict;
89    type Output = Bound<'py, Self::Target>;
90    type Error = PyErr;
91
92    #[cfg(feature = "experimental-inspect")]
93    const OUTPUT_TYPE: PyStaticExpr =
94        type_hint_subscript!(PyDict::TYPE_HINT, <&K>::OUTPUT_TYPE, <&V>::OUTPUT_TYPE);
95
96    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
97        let dict = PyDict::new(py);
98        for (k, v) in self {
99            dict.set_item(k, v)?;
100        }
101        Ok(dict)
102    }
103}
104
105impl<'py, K, V, S> FromPyObject<'_, 'py> for collections::HashMap<K, V, S>
106where
107    K: FromPyObjectOwned<'py> + cmp::Eq + hash::Hash,
108    V: FromPyObjectOwned<'py>,
109    S: hash::BuildHasher + Default,
110{
111    type Error = PyErr;
112
113    #[cfg(feature = "experimental-inspect")]
114    const INPUT_TYPE: PyStaticExpr =
115        type_hint_subscript!(&PyDict::TYPE_HINT, K::INPUT_TYPE, V::INPUT_TYPE);
116
117    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<Self, Self::Error> {
118        let dict = ob.cast::<PyDict>()?;
119        let mut ret = collections::HashMap::with_capacity_and_hasher(dict.len(), S::default());
120        for (k, v) in dict.iter() {
121            ret.insert(
122                k.extract().map_err(Into::into)?,
123                v.extract().map_err(Into::into)?,
124            );
125        }
126        Ok(ret)
127    }
128}
129
130impl<'py, K, V> FromPyObject<'_, 'py> for collections::BTreeMap<K, V>
131where
132    K: FromPyObjectOwned<'py> + cmp::Ord,
133    V: FromPyObjectOwned<'py>,
134{
135    type Error = PyErr;
136
137    #[cfg(feature = "experimental-inspect")]
138    const INPUT_TYPE: PyStaticExpr =
139        type_hint_subscript!(PyDict::TYPE_HINT, K::INPUT_TYPE, V::INPUT_TYPE);
140
141    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<Self, PyErr> {
142        let dict = ob.cast::<PyDict>()?;
143        let mut ret = collections::BTreeMap::new();
144        for (k, v) in dict.iter() {
145            ret.insert(
146                k.extract().map_err(Into::into)?,
147                v.extract().map_err(Into::into)?,
148            );
149        }
150        Ok(ret)
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use std::collections::{BTreeMap, HashMap};
158
159    #[test]
160    fn test_hashmap_to_python() {
161        Python::attach(|py| {
162            let mut map = HashMap::<i32, i32>::new();
163            map.insert(1, 1);
164
165            let py_map = (&map).into_pyobject(py).unwrap();
166
167            assert_eq!(py_map.len(), 1);
168            assert!(
169                py_map
170                    .get_item(1)
171                    .unwrap()
172                    .unwrap()
173                    .extract::<i32>()
174                    .unwrap()
175                    == 1
176            );
177            assert_eq!(map, py_map.extract().unwrap());
178        });
179    }
180
181    #[test]
182    fn test_btreemap_to_python() {
183        Python::attach(|py| {
184            let mut map = BTreeMap::<i32, i32>::new();
185            map.insert(1, 1);
186
187            let py_map = (&map).into_pyobject(py).unwrap();
188
189            assert_eq!(py_map.len(), 1);
190            assert!(
191                py_map
192                    .get_item(1)
193                    .unwrap()
194                    .unwrap()
195                    .extract::<i32>()
196                    .unwrap()
197                    == 1
198            );
199            assert_eq!(map, py_map.extract().unwrap());
200        });
201    }
202
203    #[test]
204    fn test_hashmap_into_python() {
205        Python::attach(|py| {
206            let mut map = HashMap::<i32, i32>::new();
207            map.insert(1, 1);
208
209            let py_map = map.into_pyobject(py).unwrap();
210
211            assert_eq!(py_map.len(), 1);
212            assert!(
213                py_map
214                    .get_item(1)
215                    .unwrap()
216                    .unwrap()
217                    .extract::<i32>()
218                    .unwrap()
219                    == 1
220            );
221        });
222    }
223
224    #[test]
225    fn test_btreemap_into_py() {
226        Python::attach(|py| {
227            let mut map = BTreeMap::<i32, i32>::new();
228            map.insert(1, 1);
229
230            let py_map = map.into_pyobject(py).unwrap();
231
232            assert_eq!(py_map.len(), 1);
233            assert!(
234                py_map
235                    .get_item(1)
236                    .unwrap()
237                    .unwrap()
238                    .extract::<i32>()
239                    .unwrap()
240                    == 1
241            );
242        });
243    }
244}