Module pyo3::indexmap

source ·
Available on crate feature indexmap only.
Expand description

Conversions to and from indexmap’s IndexMap.

indexmap::IndexMap is a hash table that is closely compatible with the standard std::collections::HashMap, with the difference that it preserves the insertion order when iterating over keys. It was inspired by Python’s 3.6+ dict implementation.

Dictionary order is guaranteed to be insertion order in Python, hence IndexMap is a good candidate for maintaining an equivalent behaviour in Rust.

§Setup

To use this feature, add this to your Cargo.toml:

[dependencies]
# change * to the latest versions
indexmap = "*"
pyo3 = { version = "0.21.2", features = ["indexmap"] }

Note that you must use compatible versions of indexmap and PyO3. The required indexmap version may vary based on the version of PyO3.

§Examples

Using indexmap to return a dictionary with some statistics about a list of numbers. Because of the insertion order guarantees, the Python code will always print the same result, matching users’ expectations about Python’s dict.

use indexmap::{indexmap, IndexMap};
use pyo3::prelude::*;

fn median(data: &Vec<i32>) -> f32 {
    let sorted_data = data.clone().sort();
    let mid = data.len() / 2;
    if data.len() % 2 == 0 {
        data[mid] as f32
    }
    else {
        (data[mid] + data[mid - 1]) as f32 / 2.0
    }
}

fn mean(data: &Vec<i32>) -> f32 {
    data.iter().sum::<i32>() as f32 / data.len() as f32
}
fn mode(data: &Vec<i32>) -> f32 {
    let mut frequency = IndexMap::new(); // we can use IndexMap as any hash table

    for &element in data {
        *frequency.entry(element).or_insert(0) += 1;
    }

    frequency
        .iter()
        .max_by(|a, b| a.1.cmp(&b.1))
        .map(|(k, _v)| *k)
        .unwrap() as f32
  }

#[pyfunction]
fn calculate_statistics(data: Vec<i32>) -> IndexMap<&'static str, f32> {
    indexmap! {
       "median" => median(&data),
       "mean" => mean(&data),
       "mode" => mode(&data),
    }
}

#[pymodule]
fn my_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(calculate_statistics, m)?)?;
    Ok(())
}

Python code:

from my_module import calculate_statistics

data = [1, 1, 1, 3, 4, 5]
print(calculate_statistics(data))
# always prints {"median": 2.0, "mean": 2.5, "mode": 1.0} in the same order
# if another hash table was used, the order could be random