Skip to content

Commit

Permalink
Add optional support for conversion from Hashbrown types
Browse files Browse the repository at this point in the history
This commit adds optional support for conversion from hashbrown's [1]
HashMap [2] and HashSet [3] types. The HashMap and HashSet implementation
in std::collections is a copy from HashBrown, but Hashbrown still
provides some features over the std::collections version. Primarily this
is rayon support and also using a default hasher which is faster
(although not DOS resistent). The hashbrown versions provide a drop in
replacement over std::collections to get these features. To take
advantage of native type conversion in PyO3 this commit adds hashbrown
as an optional dependency and when the feature is enabled the traits for
going between python and hashbrown::HashMap and hashbrown::HashSet are
available. This is handy for users of hashbrown which have to inline
these conversions manually in functions that take dicts as args.

[1] https://github.com/rust-lang/hashbrown
[2] https://docs.rs/hashbrown/0.8.2/hashbrown/struct.HashMap.html
[3] https://docs.rs/hashbrown/0.8.2/hashbrown/struct.HashSet.html
  • Loading branch information
mtreinish committed Aug 25, 2020
1 parent 9d73e0b commit dfa917c
Show file tree
Hide file tree
Showing 3 changed files with 166 additions and 0 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ num-complex = { version = "0.3", optional = true }
paste = { version = "0.1.6", optional = true }
pyo3cls = { path = "pyo3cls", version = "=0.11.1", optional = true }
unindent = { version = "0.1.4", optional = true }
hashbrown = { version = "0.8", optional = true }

[dev-dependencies]
assert_approx_eq = "1.1.0"
Expand Down
91 changes: 91 additions & 0 deletions src/types/dict.rs
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,97 @@ where
}
}

#[cfg(feature = "hashbrown")]
mod hashbrown_hashmap_conversion {
use super::*;
use crate::{FromPyObject, PyErr, PyObject, ToPyObject};

impl<K, V, H> ToPyObject for hashbrown::HashMap<K, V, H>
where
K: hash::Hash + cmp::Eq + ToPyObject,
V: ToPyObject,
H: hash::BuildHasher,
{
fn to_object(&self, py: Python) -> PyObject {
IntoPyDict::into_py_dict(self, py).into()
}
}

impl<K, V, H> IntoPy<PyObject> for hashbrown::HashMap<K, V, H>
where
K: hash::Hash + cmp::Eq + IntoPy<PyObject>,
V: IntoPy<PyObject>,
H: hash::BuildHasher,
{
fn into_py(self, py: Python) -> PyObject {
let iter = self
.into_iter()
.map(|(k, v)| (k.into_py(py), v.into_py(py)));
IntoPyDict::into_py_dict(iter, py).into()
}
}

impl<'source, K, V, S> FromPyObject<'source> for hashbrown::HashMap<K, V, S>
where
K: FromPyObject<'source> + cmp::Eq + hash::Hash,
V: FromPyObject<'source>,
S: hash::BuildHasher + Default,
{
fn extract(ob: &'source PyAny) -> Result<Self, PyErr> {
let dict = <PyDict as PyTryFrom>::try_from(ob)?;
let mut ret = hashbrown::HashMap::default();
for (k, v) in dict.iter() {
ret.insert(K::extract(k)?, V::extract(v)?);
}
Ok(ret)
}
}

#[test]
fn test_hashbrown_hashmap_to_python() {
let gil = Python::acquire_gil();
let py = gil.python();

let mut map = hashbrown::HashMap::<i32, i32>::new();
map.insert(1, 1);

let m = map.to_object(py);
let py_map = <PyDict as PyTryFrom>::try_from(m.as_ref(py)).unwrap();

assert!(py_map.len() == 1);
assert!(py_map.get_item(1).unwrap().extract::<i32>().unwrap() == 1);
assert_eq!(map, py_map.extract().unwrap());
}
#[test]
fn test_hashbrown_hashmap_into_python() {
let gil = Python::acquire_gil();
let py = gil.python();

let mut map = hashbrown::HashMap::<i32, i32>::new();
map.insert(1, 1);

let m: PyObject = map.into_py(py);
let py_map = <PyDict as PyTryFrom>::try_from(m.as_ref(py)).unwrap();

assert!(py_map.len() == 1);
assert!(py_map.get_item(1).unwrap().extract::<i32>().unwrap() == 1);
}

#[test]
fn test_hashbrown_hashmap_into_dict() {
let gil = Python::acquire_gil();
let py = gil.python();

let mut map = hashbrown::HashMap::<i32, i32>::new();
map.insert(1, 1);

let py_map = map.into_py_dict(py);

assert_eq!(py_map.len(), 1);
assert_eq!(py_map.get_item(1).unwrap().extract::<i32>().unwrap(), 1);
}
}

#[cfg(test)]
mod test {
use crate::conversion::IntoPy;
Expand Down
74 changes: 74 additions & 0 deletions src/types/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,80 @@ impl<'a> std::iter::IntoIterator for &'a PyFrozenSet {
}
}

#[cfg(feature = "hashbrown")]
mod hashbrown_hashset_conversion {
use super::*;
use crate::{FromPyObject, PyObject, PyResult, ToPyObject};

impl<T> ToPyObject for hashbrown::HashSet<T>
where
T: hash::Hash + Eq + ToPyObject,
{
fn to_object(&self, py: Python) -> PyObject {
let set = PySet::new::<T>(py, &[]).expect("Failed to construct empty set");
{
for val in self {
set.add(val).expect("Failed to add to set");
}
}
set.into()
}
}

impl<K, S> IntoPy<PyObject> for hashbrown::HashSet<K, S>
where
K: IntoPy<PyObject> + Eq + hash::Hash,
S: hash::BuildHasher + Default,
{
fn into_py(self, py: Python) -> PyObject {
let set = PySet::empty(py).expect("Failed to construct empty set");
{
for val in self {
set.add(val.into_py(py)).expect("Failed to add to set");
}
}
set.into()
}
}

impl<'source, K, S> FromPyObject<'source> for hashbrown::HashSet<K, S>
where
K: FromPyObject<'source> + cmp::Eq + hash::Hash,
S: hash::BuildHasher + Default,
{
fn extract(ob: &'source PyAny) -> PyResult<Self> {
let set: &PySet = ob.downcast()?;
set.iter().map(K::extract).collect()
}
}

#[test]
fn test_extract_hashbrown_hashset() {
use std::iter::FromIterator;
let gil = Python::acquire_gil();
let py = gil.python();

let set = PySet::new(py, &[1, 2, 3, 4, 5]).unwrap();
let hash_set: hashbrown::HashSet<usize> = set.extract().unwrap();
assert_eq!(
hash_set,
hashbrown::HashSet::from_iter([1, 2, 3, 4, 5].iter().copied())
);
}

#[test]
fn test_hashbrown_hashset_into_py() {
let gil = Python::acquire_gil();
let py = gil.python();

let hs: hashbrown::HashSet<u64> = [1, 2, 3, 4, 5].iter().cloned().collect();

let hso: PyObject = hs.clone().into_py(py);

assert_eq!(hs, hso.extract(py).unwrap());
}
}

#[cfg(test)]
mod test {
use super::{PyFrozenSet, PySet};
Expand Down

0 comments on commit dfa917c

Please sign in to comment.