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