Skip to main content

pyo3/conversions/
indexmap.rs

1#![cfg(feature = "indexmap")]
2
3//!  Conversions to and from [indexmap](https://docs.rs/indexmap/)’s
4//! `IndexMap`.
5//!
6//! [`indexmap::IndexMap`] is a hash table that is closely compatible with the standard [`std::collections::HashMap`],
7//! with the difference that it preserves the insertion order when iterating over keys. It was inspired
8//! by Python's 3.6+ dict implementation.
9//!
10//! Dictionary order is guaranteed to be insertion order in Python, hence IndexMap is a good candidate
11//! for maintaining an equivalent behaviour in Rust.
12//!
13//! # Setup
14//!
15//! To use this feature, add this to your **`Cargo.toml`**:
16//!
17//! ```toml
18//! [dependencies]
19//! # change * to the latest versions
20//! indexmap = "*"
21#![doc = concat!("pyo3 = { version = \"", env!("CARGO_PKG_VERSION"),  "\", features = [\"indexmap\"] }")]
22//! ```
23//!
24//! Note that you must use compatible versions of indexmap and PyO3.
25//! The required indexmap version may vary based on the version of PyO3.
26//!
27//! # Examples
28//!
29//! Using [indexmap](https://docs.rs/indexmap) to return a dictionary with some statistics
30//! about a list of numbers. Because of the insertion order guarantees, the Python code will
31//! always print the same result, matching users' expectations about Python's dict.
32//! ```rust
33//! use indexmap::{indexmap, IndexMap};
34//! use pyo3::prelude::*;
35//!
36//! fn median(data: &Vec<i32>) -> f32 {
37//!     let sorted_data = data.clone().sort();
38//!     let mid = data.len() / 2;
39//!     if data.len() % 2 == 0 {
40//!         data[mid] as f32
41//!     }
42//!     else {
43//!         (data[mid] + data[mid - 1]) as f32 / 2.0
44//!     }
45//! }
46//!
47//! fn mean(data: &Vec<i32>) -> f32 {
48//!     data.iter().sum::<i32>() as f32 / data.len() as f32
49//! }
50//! fn mode(data: &Vec<i32>) -> f32 {
51//!     let mut frequency = IndexMap::new(); // we can use IndexMap as any hash table
52//!
53//!     for &element in data {
54//!         *frequency.entry(element).or_insert(0) += 1;
55//!     }
56//!
57//!     frequency
58//!         .iter()
59//!         .max_by(|a, b| a.1.cmp(&b.1))
60//!         .map(|(k, _v)| *k)
61//!         .unwrap() as f32
62//!   }
63//!
64//! #[pyfunction]
65//! fn calculate_statistics(data: Vec<i32>) -> IndexMap<&'static str, f32> {
66//!     indexmap! {
67//!        "median" => median(&data),
68//!        "mean" => mean(&data),
69//!        "mode" => mode(&data),
70//!     }
71//! }
72//!
73//! #[pymodule]
74//! fn my_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
75//!     m.add_function(wrap_pyfunction!(calculate_statistics, m)?)?;
76//!     Ok(())
77//! }
78//! ```
79//!
80//! Python code:
81//! ```python
82//! from my_module import calculate_statistics
83//!
84//! data = [1, 1, 1, 3, 4, 5]
85//! print(calculate_statistics(data))
86//! # always prints {"median": 2.0, "mean": 2.5, "mode": 1.0} in the same order
87//! # if another hash table was used, the order could be random
88//! ```
89
90use crate::conversion::{FromPyObjectOwned, IntoPyObject};
91#[cfg(feature = "experimental-inspect")]
92use crate::inspect::PyStaticExpr;
93use crate::types::*;
94#[cfg(feature = "experimental-inspect")]
95use crate::{type_hint_subscript, PyTypeInfo};
96use crate::{Borrowed, Bound, FromPyObject, PyErr, Python};
97use std::hash;
98
99impl<'py, K, V, H> IntoPyObject<'py> for indexmap::IndexMap<K, V, H>
100where
101    K: IntoPyObject<'py> + Eq + hash::Hash,
102    V: IntoPyObject<'py>,
103    H: hash::BuildHasher,
104{
105    type Target = PyDict;
106    type Output = Bound<'py, Self::Target>;
107    type Error = PyErr;
108
109    #[cfg(feature = "experimental-inspect")]
110    const OUTPUT_TYPE: PyStaticExpr =
111        type_hint_subscript!(PyDict::TYPE_HINT, K::OUTPUT_TYPE, V::OUTPUT_TYPE);
112
113    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
114        let dict = PyDict::new(py);
115        for (k, v) in self {
116            dict.set_item(k, v)?;
117        }
118        Ok(dict)
119    }
120}
121
122impl<'a, 'py, K, V, H> IntoPyObject<'py> for &'a indexmap::IndexMap<K, V, H>
123where
124    &'a K: IntoPyObject<'py> + Eq + hash::Hash,
125    &'a V: IntoPyObject<'py>,
126    H: hash::BuildHasher,
127{
128    type Target = PyDict;
129    type Output = Bound<'py, Self::Target>;
130    type Error = PyErr;
131
132    #[cfg(feature = "experimental-inspect")]
133    const OUTPUT_TYPE: PyStaticExpr =
134        type_hint_subscript!(PyDict::TYPE_HINT, <&K>::OUTPUT_TYPE, <&V>::OUTPUT_TYPE);
135
136    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
137        let dict = PyDict::new(py);
138        for (k, v) in self {
139            dict.set_item(k, v)?;
140        }
141        Ok(dict)
142    }
143}
144
145impl<'py, K, V, S> FromPyObject<'_, 'py> for indexmap::IndexMap<K, V, S>
146where
147    K: FromPyObjectOwned<'py> + Eq + hash::Hash,
148    V: FromPyObjectOwned<'py>,
149    S: hash::BuildHasher + Default,
150{
151    type Error = PyErr;
152
153    #[cfg(feature = "experimental-inspect")]
154    const INPUT_TYPE: PyStaticExpr =
155        type_hint_subscript!(PyDict::TYPE_HINT, K::INPUT_TYPE, V::INPUT_TYPE);
156
157    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<Self, Self::Error> {
158        let dict = ob.cast::<PyDict>()?;
159        let mut ret = indexmap::IndexMap::with_capacity_and_hasher(dict.len(), S::default());
160        for (k, v) in dict.iter() {
161            ret.insert(
162                k.extract().map_err(Into::into)?,
163                v.extract().map_err(Into::into)?,
164            );
165        }
166        Ok(ret)
167    }
168}
169
170#[cfg(test)]
171mod test_indexmap {
172
173    use crate::types::*;
174    use crate::{IntoPyObject, Python};
175
176    #[test]
177    fn test_indexmap_indexmap_into_pyobject() {
178        Python::attach(|py| {
179            let mut map = indexmap::IndexMap::<i32, i32>::new();
180            map.insert(1, 1);
181
182            let py_map = (&map).into_pyobject(py).unwrap();
183
184            assert_eq!(py_map.len(), 1);
185            assert!(
186                py_map
187                    .get_item(1)
188                    .unwrap()
189                    .unwrap()
190                    .extract::<i32>()
191                    .unwrap()
192                    == 1
193            );
194            assert_eq!(
195                map,
196                py_map.extract::<indexmap::IndexMap::<i32, i32>>().unwrap()
197            );
198        });
199    }
200
201    #[test]
202    fn test_indexmap_indexmap_into_dict() {
203        Python::attach(|py| {
204            let mut map = indexmap::IndexMap::<i32, i32>::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_indexmap_indexmap_insertion_order_round_trip() {
224        Python::attach(|py| {
225            let n = 20;
226            let mut map = indexmap::IndexMap::<i32, i32>::new();
227
228            for i in 1..=n {
229                if i % 2 == 1 {
230                    map.insert(i, i);
231                } else {
232                    map.insert(n - i, i);
233                }
234            }
235
236            let py_map = (&map).into_py_dict(py).unwrap();
237
238            let trip_map = py_map.extract::<indexmap::IndexMap<i32, i32>>().unwrap();
239
240            for (((k1, v1), (k2, v2)), (k3, v3)) in
241                map.iter().zip(py_map.iter()).zip(trip_map.iter())
242            {
243                let k2 = k2.extract::<i32>().unwrap();
244                let v2 = v2.extract::<i32>().unwrap();
245                assert_eq!((k1, v1), (&k2, &v2));
246                assert_eq!((k1, v1), (k3, v3));
247                assert_eq!((&k2, &v2), (k3, v3));
248            }
249        });
250    }
251}
⚠️ Internal Docs ⚠️ Not Public API 👉 Official Docs Here