From dfeae473e5cdd7c23a41d6a0e2a0796e26a087ba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Orhun=20Parmaks=C4=B1z?= Date: Wed, 11 Oct 2023 13:44:43 +0300 Subject: [PATCH] Add support for `SmallVec` in conversion traits (#3440) --- Cargo.toml | 2 + newsfragments/3507.added.md | 1 + src/conversions/mod.rs | 1 + src/conversions/smallvec.rs | 140 ++++++++++++++++++++++++++++++++++++ 4 files changed, 144 insertions(+) create mode 100644 newsfragments/3507.added.md create mode 100644 src/conversions/smallvec.rs diff --git a/Cargo.toml b/Cargo.toml index 28716777..abf469eb 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 00000000..2068ab4c --- /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 5544dc23..a9c2b0cd 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 00000000..d2e84421 --- /dev/null +++ b/src/conversions/smallvec.rs @@ -0,0 +1,140 @@ +#![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::Item::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::{PyDict, 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]); + }); + } + + #[test] + fn test_smallvec_from_py_object_fails() { + Python::with_gil(|py| { + let dict = PyDict::new(py); + let sv: PyResult> = dict.extract(); + assert_eq!( + sv.unwrap_err().to_string(), + "TypeError: 'dict' object cannot be converted to 'Sequence'" + ); + }); + } + + #[test] + fn test_smallvec_to_object() { + Python::with_gil(|py| { + let sv: SmallVec<[u64; 8]> = [1, 2, 3, 4, 5].iter().cloned().collect(); + let hso: PyObject = sv.to_object(py); + let l = PyList::new(py, [1, 2, 3, 4, 5]); + assert!(l.eq(hso).unwrap()); + }); + } +}