diff --git a/Cargo.toml b/Cargo.toml index 2871677754e..abf469eb88a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,6 +41,7 @@ num-bigint = { version = "0.4", optional = true } num-complex = { version = ">= 0.2, < 0.5", optional = true } rust_decimal = { version = "1.0.0", default-features = false, optional = true } serde = { version = "1.0", optional = true } +smallvec = { version = "1.11.1", optional = true } [dev-dependencies] assert_approx_eq = "1.1.0" @@ -104,6 +105,7 @@ full = [ "num-bigint", "num-complex", "hashbrown", + "smallvec", "serde", "indexmap", "eyre", diff --git a/newsfragments/3507.added.md b/newsfragments/3507.added.md new file mode 100644 index 00000000000..2068ab4c3f7 --- /dev/null +++ b/newsfragments/3507.added.md @@ -0,0 +1 @@ +Add `smallvec` feature to add `ToPyObject`, `IntoPy` and `FromPyObject` implementations for `smallvec::SmallVec`. diff --git a/src/conversions/mod.rs b/src/conversions/mod.rs index 5544dc23532..a9c2b0cd2a6 100644 --- a/src/conversions/mod.rs +++ b/src/conversions/mod.rs @@ -9,4 +9,5 @@ pub mod num_bigint; pub mod num_complex; pub mod rust_decimal; pub mod serde; +pub mod smallvec; mod std; diff --git a/src/conversions/smallvec.rs b/src/conversions/smallvec.rs new file mode 100644 index 00000000000..fd9ec922810 --- /dev/null +++ b/src/conversions/smallvec.rs @@ -0,0 +1,118 @@ +#![cfg(feature = "smallvec")] + +//! Conversions to and from [smallvec](https://docs.rs/smallvec/). +//! +//! # Setup +//! +//! To use this feature, add this to your **`Cargo.toml`**: +//! +//! ```toml +//! [dependencies] +//! # change * to the latest versions +//! smallvec = "*" +#![doc = concat!("pyo3 = { version = \"", env!("CARGO_PKG_VERSION"), "\", features = [\"smallvec\"] }")] +//! ``` +//! +//! Note that you must use compatible versions of smallvec and PyO3. +//! The required smallvec version may vary based on the version of PyO3. +use crate::exceptions::PyTypeError; +#[cfg(feature = "experimental-inspect")] +use crate::inspect::types::TypeInfo; +use crate::types::list::new_from_iter; +use crate::types::{PySequence, PyString}; +use crate::{ + ffi, FromPyObject, IntoPy, PyAny, PyDowncastError, PyObject, PyResult, Python, ToPyObject, +}; +use smallvec::{Array, SmallVec}; + +impl ToPyObject for SmallVec +where + A: Array, + A::Item: ToPyObject, +{ + fn to_object(&self, py: Python<'_>) -> PyObject { + self.as_slice().to_object(py) + } +} + +impl IntoPy for SmallVec +where + A: Array, + A::Item: IntoPy, +{ + fn into_py(self, py: Python<'_>) -> PyObject { + let mut iter = self.into_iter().map(|e| e.into_py(py)); + let list = new_from_iter(py, &mut iter); + list.into() + } + + #[cfg(feature = "experimental-inspect")] + fn type_output() -> TypeInfo { + TypeInfo::list_of(A::Item::type_output()) + } +} + +impl<'a, A> FromPyObject<'a> for SmallVec +where + A: Array, + A::Item: FromPyObject<'a>, +{ + fn extract(obj: &'a PyAny) -> PyResult { + if obj.is_instance_of::() { + return Err(PyTypeError::new_err("Can't extract `str` to `SmallVec`")); + } + extract_sequence(obj) + } + + #[cfg(feature = "experimental-inspect")] + fn type_input() -> TypeInfo { + TypeInfo::sequence_of(A::type_input()) + } +} + +fn extract_sequence<'s, A>(obj: &'s PyAny) -> PyResult> +where + A: Array, + A::Item: FromPyObject<'s>, +{ + // Types that pass `PySequence_Check` usually implement enough of the sequence protocol + // to support this function and if not, we will only fail extraction safely. + let seq: &PySequence = unsafe { + if ffi::PySequence_Check(obj.as_ptr()) != 0 { + obj.downcast_unchecked() + } else { + return Err(PyDowncastError::new(obj, "Sequence").into()); + } + }; + + let mut sv = SmallVec::with_capacity(seq.len().unwrap_or(0)); + for item in seq.iter()? { + sv.push(item?.extract::()?); + } + Ok(sv) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::PyList; + + #[test] + fn test_smallvec_into_py() { + Python::with_gil(|py| { + let sv: SmallVec<[u64; 8]> = [1, 2, 3, 4, 5].iter().cloned().collect(); + let hso: PyObject = sv.clone().into_py(py); + let l = PyList::new(py, [1, 2, 3, 4, 5]); + assert!(l.eq(hso).unwrap()); + }); + } + + #[test] + fn test_smallvec_from_py_object() { + Python::with_gil(|py| { + let l = PyList::new(py, [1, 2, 3, 4, 5]); + let sv: SmallVec<[u64; 8]> = l.extract().unwrap(); + assert_eq!(sv.as_slice(), [1, 2, 3, 4, 5]); + }); + } +}