|
| 1 | +//! Conversions to and from [indexmap](https://docs.rs/indexmap/)’s |
| 2 | +//! `IndexMap`. |
| 3 | +//! |
| 4 | +//! [`indexmap::IndexMap`] is a hash table that is closely compatible with the standard [`std::collections::HashMap`], |
| 5 | +//! with the difference that it preserves the insertion order when iterating over keys. It was inspired |
| 6 | +//! by Python's 3.6+ dict implementation. |
| 7 | +//! |
| 8 | +//! Dictionary order is guaranteed to be insertion order in Python, hence IndexMap is a good candidate |
| 9 | +//! for maintaining an equivalent behaviour in Rust. |
| 10 | +//! |
| 11 | +//! # Setup |
| 12 | +//! |
| 13 | +//! To use this feature, add this to your **`Cargo.toml`**: |
| 14 | +//! |
| 15 | +//! ```toml |
| 16 | +//! [dependencies] |
| 17 | +//! # change * to the latest versions |
| 18 | +//! indexmap = "*" |
| 19 | +// workaround for `extended_key_value_attributes`: https://github.com/rust-lang/rust/issues/82768#issuecomment-803935643 |
| 20 | +#![cfg_attr(docsrs, cfg_attr(docsrs, doc = concat!("pyo3 = { version = \"", env!("CARGO_PKG_VERSION"), "\", features = [\"indexmap\"] }")))] |
| 21 | +#![cfg_attr( |
| 22 | + not(docsrs), |
| 23 | + doc = "pyo3 = { version = \"*\", features = [\"indexmap\"] }" |
| 24 | +)] |
| 25 | +//! ``` |
| 26 | +//! |
| 27 | +//! Note that you must use compatible versions of indexmap and PyO3. |
| 28 | +//! The required indexmap version may vary based on the version of PyO3. |
| 29 | +//! |
| 30 | +//! # Examples |
| 31 | +//! |
| 32 | +//! Using [indexmap](https://docs.rs/indexmap) to return a dictionary with some statistics |
| 33 | +//! about a list of numbers. Because of the insertion order guarantees, the Python code will |
| 34 | +//! always print the same result, matching users' expectations about Python's dict. |
| 35 | +//! |
| 36 | +//! ```rust |
| 37 | +//! use indexmap::{indexmap, IndexMap}; |
| 38 | +//! use pyo3::prelude::*; |
| 39 | +//! |
| 40 | +//! fn median(data: &Vec<i32>) -> f32 { |
| 41 | +//! let sorted_data = data.clone().sort(); |
| 42 | +//! let mid = data.len() / 2; |
| 43 | +//! if (data.len() % 2 == 0) { |
| 44 | +//! data[mid] as f32 |
| 45 | +//! } |
| 46 | +//! else { |
| 47 | +//! (data[mid] + data[mid - 1]) as f32 / 2.0 |
| 48 | +//! } |
| 49 | +//! } |
| 50 | +//! |
| 51 | +//! fn mean(data: &Vec<i32>) -> f32 { |
| 52 | +//! data.iter().sum::<i32>() as f32 / data.len() as f32 |
| 53 | +//! } |
| 54 | +//! fn mode(data: &Vec<i32>) -> f32 { |
| 55 | +//! let mut frequency = IndexMap::new(); // we can use IndexMap as any hash table |
| 56 | +//! |
| 57 | +//! for &element in data { |
| 58 | +//! *frequency.entry(element).or_insert(0) += 1; |
| 59 | +//! } |
| 60 | +//! |
| 61 | +//! frequency |
| 62 | +//! .iter() |
| 63 | +//! .max_by(|a, b| a.1.cmp(&b.1)) |
| 64 | +//! .map(|(k, _v)| *k) |
| 65 | +//! .unwrap() as f32 |
| 66 | +//! } |
| 67 | +//! |
| 68 | +//! #[pyfunction] |
| 69 | +//! fn calculate_statistics(data: Vec<i32>) -> IndexMap<&'static str, f32> { |
| 70 | +//! indexmap!{ |
| 71 | +//! "median" => median(&data), |
| 72 | +//! "mean" => mean(&data), |
| 73 | +//! "mode" => mode(&data), |
| 74 | +//! } |
| 75 | +//! } |
| 76 | +//! |
| 77 | +//! #[pymodule] |
| 78 | +//! fn my_module(_py: Python, m: &PyModule) -> PyResult<()> { |
| 79 | +//! m.add_function(wrap_pyfunction!(calculate_statistics, m)?)?; |
| 80 | +//! Ok(()) |
| 81 | +//! } |
| 82 | +//! ``` |
| 83 | +//! |
| 84 | +//! Python code: |
| 85 | +//! ```python |
| 86 | +//! from my_module import calculate_statistics |
| 87 | +//! |
| 88 | +//! data = [1, 1, 1, 3, 4, 5] |
| 89 | +//! print(calculate_statistics(data)) |
| 90 | +//! # always prints {"median": 2.0, "mean": 2.5, "mode": 1.0} in the same order |
| 91 | +//! # if another hash table was used, the order could be random |
| 92 | +//! ``` |
| 93 | +
|
| 94 | +use crate::types::*; |
| 95 | +use crate::{FromPyObject, IntoPy, PyErr, PyObject, PyTryFrom, Python, ToPyObject}; |
| 96 | +use std::{cmp, hash}; |
| 97 | + |
| 98 | +impl<K, V, H> ToPyObject for indexmap::IndexMap<K, V, H> |
| 99 | +where |
| 100 | + K: hash::Hash + cmp::Eq + ToPyObject, |
| 101 | + V: ToPyObject, |
| 102 | + H: hash::BuildHasher, |
| 103 | +{ |
| 104 | + fn to_object(&self, py: Python) -> PyObject { |
| 105 | + IntoPyDict::into_py_dict(self, py).into() |
| 106 | + } |
| 107 | +} |
| 108 | + |
| 109 | +impl<K, V, H> IntoPy<PyObject> for indexmap::IndexMap<K, V, H> |
| 110 | +where |
| 111 | + K: hash::Hash + cmp::Eq + IntoPy<PyObject>, |
| 112 | + V: IntoPy<PyObject>, |
| 113 | + H: hash::BuildHasher, |
| 114 | +{ |
| 115 | + fn into_py(self, py: Python) -> PyObject { |
| 116 | + let iter = self |
| 117 | + .into_iter() |
| 118 | + .map(|(k, v)| (k.into_py(py), v.into_py(py))); |
| 119 | + IntoPyDict::into_py_dict(iter, py).into() |
| 120 | + } |
| 121 | +} |
| 122 | + |
| 123 | +impl<'source, K, V, S> FromPyObject<'source> for indexmap::IndexMap<K, V, S> |
| 124 | +where |
| 125 | + K: FromPyObject<'source> + cmp::Eq + hash::Hash, |
| 126 | + V: FromPyObject<'source>, |
| 127 | + S: hash::BuildHasher + Default, |
| 128 | +{ |
| 129 | + fn extract(ob: &'source PyAny) -> Result<Self, PyErr> { |
| 130 | + let dict = <PyDict as PyTryFrom>::try_from(ob)?; |
| 131 | + let mut ret = indexmap::IndexMap::with_capacity_and_hasher(dict.len(), S::default()); |
| 132 | + for (k, v) in dict.iter() { |
| 133 | + ret.insert(K::extract(k)?, V::extract(v)?); |
| 134 | + } |
| 135 | + Ok(ret) |
| 136 | + } |
| 137 | +} |
| 138 | + |
| 139 | +#[cfg(test)] |
| 140 | +mod test_indexmap { |
| 141 | + |
| 142 | + use crate::types::*; |
| 143 | + use crate::{IntoPy, PyObject, PyTryFrom, Python, ToPyObject}; |
| 144 | + |
| 145 | + #[test] |
| 146 | + fn test_indexmap_indexmap_to_python() { |
| 147 | + Python::with_gil(|py| { |
| 148 | + let mut map = indexmap::IndexMap::<i32, i32>::new(); |
| 149 | + map.insert(1, 1); |
| 150 | + |
| 151 | + let m = map.to_object(py); |
| 152 | + let py_map = <PyDict as PyTryFrom>::try_from(m.as_ref(py)).unwrap(); |
| 153 | + |
| 154 | + assert!(py_map.len() == 1); |
| 155 | + assert!(py_map.get_item(1).unwrap().extract::<i32>().unwrap() == 1); |
| 156 | + assert_eq!( |
| 157 | + map, |
| 158 | + py_map.extract::<indexmap::IndexMap::<i32, i32>>().unwrap() |
| 159 | + ); |
| 160 | + }); |
| 161 | + } |
| 162 | + |
| 163 | + #[test] |
| 164 | + fn test_indexmap_indexmap_into_python() { |
| 165 | + Python::with_gil(|py| { |
| 166 | + let mut map = indexmap::IndexMap::<i32, i32>::new(); |
| 167 | + map.insert(1, 1); |
| 168 | + |
| 169 | + let m: PyObject = map.into_py(py); |
| 170 | + let py_map = <PyDict as PyTryFrom>::try_from(m.as_ref(py)).unwrap(); |
| 171 | + |
| 172 | + assert!(py_map.len() == 1); |
| 173 | + assert!(py_map.get_item(1).unwrap().extract::<i32>().unwrap() == 1); |
| 174 | + }); |
| 175 | + } |
| 176 | + |
| 177 | + #[test] |
| 178 | + fn test_indexmap_indexmap_into_dict() { |
| 179 | + Python::with_gil(|py| { |
| 180 | + let mut map = indexmap::IndexMap::<i32, i32>::new(); |
| 181 | + map.insert(1, 1); |
| 182 | + |
| 183 | + let py_map = map.into_py_dict(py); |
| 184 | + |
| 185 | + assert_eq!(py_map.len(), 1); |
| 186 | + assert_eq!(py_map.get_item(1).unwrap().extract::<i32>().unwrap(), 1); |
| 187 | + }); |
| 188 | + } |
| 189 | + |
| 190 | + #[test] |
| 191 | + fn test_indexmap_indexmap_insertion_order_round_trip() { |
| 192 | + Python::with_gil(|py| { |
| 193 | + let n = 20; |
| 194 | + let mut map = indexmap::IndexMap::<i32, i32>::new(); |
| 195 | + |
| 196 | + for i in 1..=n { |
| 197 | + if i % 2 == 1 { |
| 198 | + map.insert(i, i); |
| 199 | + } else { |
| 200 | + map.insert(n - i, i); |
| 201 | + } |
| 202 | + } |
| 203 | + |
| 204 | + let py_map = map.clone().into_py_dict(py); |
| 205 | + |
| 206 | + let trip_map = py_map.extract::<indexmap::IndexMap<i32, i32>>().unwrap(); |
| 207 | + |
| 208 | + for (((k1, v1), (k2, v2)), (k3, v3)) in |
| 209 | + map.iter().zip(py_map.iter()).zip(trip_map.iter()) |
| 210 | + { |
| 211 | + let k2 = k2.extract::<i32>().unwrap(); |
| 212 | + let v2 = v2.extract::<i32>().unwrap(); |
| 213 | + assert_eq!((k1, v1), (&k2, &v2)); |
| 214 | + assert_eq!((k1, v1), (k3, v3)); |
| 215 | + assert_eq!((&k2, &v2), (k3, v3)); |
| 216 | + } |
| 217 | + }); |
| 218 | + } |
| 219 | +} |
0 commit comments