pyo3 0.29.3

use alloc::vec::Vec;
use core::{cmp, hash};
use std::collections;

#[cfg(feature = "experimental-inspect")]
use crate::inspect::{type_hint_subscript, PyStaticExpr};
#[cfg(feature = "experimental-inspect")]
use crate::type_object::PyTypeInfo;
use crate::{
    conversion::{FromPyObjectOwned, IntoPyObject},
    types::{
        any::PyAnyMethods, frozenset::PyFrozenSetMethods, set::PySetMethods, PyFrozenSet, PySet,
    },
    Borrowed, Bound, FromPyObject, PyAny, PyErr, Python,
};

impl<'py, K, S> IntoPyObject<'py> for collections::HashSet<K, S>
where
    K: IntoPyObject<'py> + Eq + hash::Hash,
    S: hash::BuildHasher + Default,
{
    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 collections::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 collections::HashSet<K, S>
where
    K: FromPyObjectOwned<'py> + cmp::Eq + hash::Hash,
    S: hash::BuildHasher + Default,
{
    type Error = PyErr;

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

    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<Self, Self::Error> {
        match ob.cast::<PySet>() {
            Ok(set) => {
                let mut result = Self::with_capacity_and_hasher(set.len(), S::default());
                for item in set.iter() {
                    result.insert(item.extract().map_err(Into::into)?);
                }
                Ok(result)
            }
            Err(err) => {
                if let Ok(frozen_set) = ob.cast::<PyFrozenSet>() {
                    let mut result = Self::with_capacity_and_hasher(frozen_set.len(), S::default());
                    for item in frozen_set.iter() {
                        result.insert(item.extract().map_err(Into::into)?);
                    }
                    Ok(result)
                } else {
                    Err(PyErr::from(err))
                }
            }
        }
    }
}

impl<'py, K> IntoPyObject<'py> for collections::BTreeSet<K>
where
    K: IntoPyObject<'py> + cmp::Ord,
{
    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> IntoPyObject<'py> for &'a collections::BTreeSet<K>
where
    &'a K: IntoPyObject<'py> + cmp::Ord,
    K: 'a,
{
    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> FromPyObject<'_, 'py> for collections::BTreeSet<K>
where
    K: FromPyObjectOwned<'py> + cmp::Ord,
{
    type Error = PyErr;

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

    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result<Self, Self::Error> {
        match ob.cast::<PySet>() {
            Ok(set) => {
                let mut values = Vec::with_capacity(set.len());
                for item in set.iter() {
                    values.push(item.extract().map_err(Into::into)?);
                }
                Ok(values.into_iter().collect())
            }
            Err(err) => {
                if let Ok(frozen_set) = ob.cast::<PyFrozenSet>() {
                    let mut values = Vec::with_capacity(frozen_set.len());
                    for item in frozen_set.iter() {
                        values.push(item.extract().map_err(Into::into)?);
                    }
                    Ok(values.into_iter().collect())
                } else {
                    Err(PyErr::from(err))
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::types::{any::PyAnyMethods, PyFrozenSet, PySet};
    use crate::{IntoPyObject, Python};
    use alloc::collections::BTreeSet;
    use std::collections::HashSet;

    #[test]
    fn test_extract_hashset() {
        Python::attach(|py| {
            let set = PySet::new(py, [1, 2, 3, 4, 5]).unwrap();
            let hash_set: HashSet<usize> = 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: HashSet<usize> = set.extract().unwrap();
            assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
        });
    }

    #[test]
    fn test_extract_btreeset() {
        Python::attach(|py| {
            let set = PySet::new(py, [1, 2, 3, 4, 5]).unwrap();
            let hash_set: BTreeSet<usize> = 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: BTreeSet<usize> = set.extract().unwrap();
            assert_eq!(hash_set, [1, 2, 3, 4, 5].iter().copied().collect());
        });
    }

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

            let bto = (&bt).into_pyobject(py).unwrap();
            let hso = (&hs).into_pyobject(py).unwrap();

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