pyforge 0.3.0

#![cfg(feature = "hashbrown")]

//!  Conversions to and from [hashbrown](https://docs.rs/hashbrown/)’s
//! `HashMap` and `HashSet`.
//!
//! # Setup
//!
//! To use this feature, add this to your **`Cargo.toml`**:
//!
//! ```toml
//! [dependencies]
//! # change * to the latest versions
//! hashbrown = "*"
#![doc = concat!("pyo3 = { version = \"", env!("CARGO_PKG_VERSION"),  "\", features = [\"hashbrown\"] }")]
//! ```
//!
//! Note that you must use compatible versions of hashbrown and PyForge.
//! The required hashbrown version may vary based on the version of PyForge.
#[cfg(feature = "experimental-inspect")]
use crate::inspect::PyStaticExpr;
use crate::{
    conversion::{FromPyObjectOwned, IntoPyObject},
    types::{
        any::PyAnyMethods, dict::PyDictMethods, frozenset::PyFrozenSetMethods, set::PySetMethods,
        PyDict, PyFrozenSet, PySet,
    },
    Borrowed, Bound, FromPyObject, PyAny, PyErr, PyResult, Python,
};
#[cfg(feature = "experimental-inspect")]
use crate::{type_hint_subscript, type_hint_union, PyTypeInfo};
use std::hash;

impl<'py, K, V, H> IntoPyObject<'py> for hashbrown::HashMap<K, V, H>
where
    K: IntoPyObject<'py> + Eq + hash::Hash,
    V: IntoPyObject<'py>,
    H: hash::BuildHasher,
{
    type Target = PyDict;
    type Output = Bound<'py, Self::Target>;
    type Error = PyErr;

    #[cfg(feature = "experimental-inspect")]
    const OUTPUT_TYPE: PyStaticExpr =
        type_hint_subscript!(PyDict::TYPE_HINT, K::OUTPUT_TYPE, V::OUTPUT_TYPE);

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        let dict = PyDict::new(py);
        for (k, v) in self {
            dict.set_item(k, v)?;
        }
        Ok(dict)
    }
}

impl<'a, 'py, K, V, H> IntoPyObject<'py> for &'a hashbrown::HashMap<K, V, H>
where
    &'a K: IntoPyObject<'py> + Eq + hash::Hash,
    &'a V: IntoPyObject<'py>,
    H: hash::BuildHasher,
{
    type Target = PyDict;
    type Output = Bound<'py, Self::Target>;
    type Error = PyErr;

    #[cfg(feature = "experimental-inspect")]
    const OUTPUT_TYPE: PyStaticExpr =
        type_hint_subscript!(PyDict::TYPE_HINT, <&K>::OUTPUT_TYPE, <&V>::OUTPUT_TYPE);

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        let dict = PyDict::new(py);
        for (k, v) in self {
            dict.set_item(k, v)?;
        }
        Ok(dict)
    }
}

impl<'py, K, V, S> FromPyObject<'_, 'py> for hashbrown::HashMap<K, V, S>
where
    K: FromPyObjectOwned<'py> + Eq + hash::Hash,
    V: FromPyObjectOwned<'py>,
    S: hash::BuildHasher + Default,
{
    type Error = PyErr;

    #[cfg(feature = "experimental-inspect")]
    const INPUT_TYPE: PyStaticExpr =
        type_hint_subscript!(PyDict::TYPE_HINT, K::INPUT_TYPE, V::INPUT_TYPE);

    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<Self, PyErr> {
        let dict = ob.cast::<PyDict>()?;
        let mut ret = hashbrown::HashMap::with_capacity_and_hasher(dict.len(), S::default());
        for (k, v) in dict.iter() {
            ret.insert(
                k.extract().map_err(Into::into)?,
                v.extract().map_err(Into::into)?,
            );
        }
        Ok(ret)
    }
}

impl<'py, K, H> IntoPyObject<'py> for hashbrown::HashSet<K, H>
where
    K: IntoPyObject<'py> + Eq + hash::Hash,
    H: hash::BuildHasher,
{
    type Target = PySet;
    type Output = Bound<'py, Self::Target>;
    type Error = PyErr;

    #[cfg(feature = "experimental-inspect")]
    const OUTPUT_TYPE: PyStaticExpr = type_hint_subscript!(PySet::TYPE_HINT, K::OUTPUT_TYPE);

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        PySet::new(py, self)
    }
}

impl<'a, 'py, K, H> IntoPyObject<'py> for &'a hashbrown::HashSet<K, H>
where
    &'a K: IntoPyObject<'py> + Eq + hash::Hash,
    H: hash::BuildHasher,
{
    type Target = PySet;
    type Output = Bound<'py, Self::Target>;
    type Error = PyErr;

    #[cfg(feature = "experimental-inspect")]
    const OUTPUT_TYPE: PyStaticExpr = type_hint_subscript!(PySet::TYPE_HINT, <&K>::OUTPUT_TYPE);

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        PySet::new(py, self)
    }
}

impl<'py, K, S> FromPyObject<'_, 'py> for hashbrown::HashSet<K, S>
where
    K: FromPyObjectOwned<'py> + Eq + hash::Hash,
    S: hash::BuildHasher + Default,
{
    type Error = PyErr;

    #[cfg(feature = "experimental-inspect")]
    const INPUT_TYPE: PyStaticExpr = type_hint_union!(
        type_hint_subscript!(PySet::TYPE_HINT, K::INPUT_TYPE),
        type_hint_subscript!(PyFrozenSet::TYPE_HINT, K::INPUT_TYPE)
    );

    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {
        match ob.cast::<PySet>() {
            Ok(set) => set
                .iter()
                .map(|any| any.extract().map_err(Into::into))
                .collect(),
            Err(err) => {
                if let Ok(frozen_set) = ob.cast::<PyFrozenSet>() {
                    frozen_set
                        .iter()
                        .map(|any| any.extract().map_err(Into::into))
                        .collect()
                } else {
                    Err(PyErr::from(err))
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::IntoPyDict;
    use std::collections::hash_map::RandomState;

    #[test]
    fn test_hashbrown_hashmap_into_pyobject() {
        Python::attach(|py| {
            let mut map =
                hashbrown::HashMap::<i32, i32, RandomState>::with_hasher(RandomState::new());
            map.insert(1, 1);

            let py_map = (&map).into_pyobject(py).unwrap();

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

    #[test]
    fn test_hashbrown_hashmap_into_dict() {
        Python::attach(|py| {
            let mut map =
                hashbrown::HashMap::<i32, i32, RandomState>::with_hasher(RandomState::new());
            map.insert(1, 1);

            let py_map = map.into_py_dict(py).unwrap();

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

    #[test]
    fn test_extract_hashbrown_hashset() {
        Python::attach(|py| {
            let set = PySet::new(py, [1, 2, 3, 4, 5]).unwrap();
            let hash_set: hashbrown::HashSet<usize, RandomState> = set.extract().unwrap();
            assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());

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

    #[test]
    fn test_hashbrown_hashset_into_pyobject() {
        Python::attach(|py| {
            let hs: hashbrown::HashSet<u64, RandomState> =
                [1, 2, 3, 4, 5].iter().cloned().collect();

            let hso = hs.clone().into_pyobject(py).unwrap();

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