pyo3/conversions/
smallvec.rs

1#![cfg(feature = "smallvec")]
2
3//!  Conversions to and from [smallvec](https://docs.rs/smallvec/).
4//!
5//! # Setup
6//!
7//! To use this feature, add this to your **`Cargo.toml`**:
8//!
9//! ```toml
10//! [dependencies]
11//! # change * to the latest versions
12//! smallvec = "*"
13#![doc = concat!("pyo3 = { version = \"", env!("CARGO_PKG_VERSION"),  "\", features = [\"smallvec\"] }")]
14//! ```
15//!
16//! Note that you must use compatible versions of smallvec and PyO3.
17//! The required smallvec version may vary based on the version of PyO3.
18use crate::conversion::{FromPyObjectOwned, IntoPyObject};
19use crate::exceptions::PyTypeError;
20#[cfg(feature = "experimental-inspect")]
21use crate::inspect::types::TypeInfo;
22#[cfg(feature = "experimental-inspect")]
23use crate::inspect::TypeHint;
24use crate::types::any::PyAnyMethods;
25use crate::types::{PySequence, PyString};
26use crate::{
27    err::CastError, ffi, Borrowed, Bound, FromPyObject, PyAny, PyErr, PyResult, PyTypeInfo, Python,
28};
29use smallvec::{Array, SmallVec};
30
31impl<'py, A> IntoPyObject<'py> for SmallVec<A>
32where
33    A: Array,
34    A::Item: IntoPyObject<'py>,
35{
36    type Target = PyAny;
37    type Output = Bound<'py, Self::Target>;
38    type Error = PyErr;
39
40    #[cfg(feature = "experimental-inspect")]
41    const OUTPUT_TYPE: TypeHint = A::Item::SEQUENCE_OUTPUT_TYPE;
42
43    /// Turns [`SmallVec<u8>`] into [`PyBytes`], all other `T`s will be turned into a [`PyList`]
44    ///
45    /// [`PyBytes`]: crate::types::PyBytes
46    /// [`PyList`]: crate::types::PyList
47    #[inline]
48    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
49        <A::Item>::owned_sequence_into_pyobject(self, py, crate::conversion::private::Token)
50    }
51
52    #[cfg(feature = "experimental-inspect")]
53    fn type_output() -> TypeInfo {
54        TypeInfo::list_of(A::Item::type_output())
55    }
56}
57
58impl<'a, 'py, A> IntoPyObject<'py> for &'a SmallVec<A>
59where
60    A: Array,
61    &'a A::Item: IntoPyObject<'py>,
62{
63    type Target = PyAny;
64    type Output = Bound<'py, Self::Target>;
65    type Error = PyErr;
66
67    #[inline]
68    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
69        self.as_slice().into_pyobject(py)
70    }
71
72    #[cfg(feature = "experimental-inspect")]
73    fn type_output() -> TypeInfo {
74        TypeInfo::list_of(<&A::Item>::type_output())
75    }
76}
77
78impl<'py, A> FromPyObject<'_, 'py> for SmallVec<A>
79where
80    A: Array,
81    A::Item: FromPyObjectOwned<'py>,
82{
83    type Error = PyErr;
84
85    fn extract(obj: Borrowed<'_, 'py, PyAny>) -> Result<Self, Self::Error> {
86        if obj.is_instance_of::<PyString>() {
87            return Err(PyTypeError::new_err("Can't extract `str` to `SmallVec`"));
88        }
89        extract_sequence(obj)
90    }
91
92    #[cfg(feature = "experimental-inspect")]
93    fn type_input() -> TypeInfo {
94        TypeInfo::sequence_of(A::Item::type_input())
95    }
96}
97
98fn extract_sequence<'py, A>(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<SmallVec<A>>
99where
100    A: Array,
101    A::Item: FromPyObjectOwned<'py>,
102{
103    // Types that pass `PySequence_Check` usually implement enough of the sequence protocol
104    // to support this function and if not, we will only fail extraction safely.
105    let seq = unsafe {
106        if ffi::PySequence_Check(obj.as_ptr()) != 0 {
107            obj.cast_unchecked::<PySequence>()
108        } else {
109            return Err(CastError::new(obj, PySequence::type_object(obj.py()).into_any()).into());
110        }
111    };
112
113    let mut sv = SmallVec::with_capacity(seq.len().unwrap_or(0));
114    for item in seq.try_iter()? {
115        sv.push(item?.extract::<A::Item>().map_err(Into::into)?);
116    }
117    Ok(sv)
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::types::{PyBytes, PyBytesMethods, PyDict, PyList};
124
125    #[test]
126    fn test_smallvec_from_py_object() {
127        Python::attach(|py| {
128            let l = PyList::new(py, [1, 2, 3, 4, 5]).unwrap();
129            let sv: SmallVec<[u64; 8]> = l.extract().unwrap();
130            assert_eq!(sv.as_slice(), [1, 2, 3, 4, 5]);
131        });
132    }
133
134    #[test]
135    fn test_smallvec_from_py_object_fails() {
136        Python::attach(|py| {
137            let dict = PyDict::new(py);
138            let sv: PyResult<SmallVec<[u64; 8]>> = dict.extract();
139            assert_eq!(
140                sv.unwrap_err().to_string(),
141                "TypeError: 'dict' object cannot be cast as 'Sequence'"
142            );
143        });
144    }
145
146    #[test]
147    fn test_smallvec_into_pyobject() {
148        Python::attach(|py| {
149            let sv: SmallVec<[u64; 8]> = [1, 2, 3, 4, 5].iter().cloned().collect();
150            let hso = sv.into_pyobject(py).unwrap();
151            let l = PyList::new(py, [1, 2, 3, 4, 5]).unwrap();
152            assert!(l.eq(hso).unwrap());
153        });
154    }
155
156    #[test]
157    fn test_smallvec_intopyobject_impl() {
158        Python::attach(|py| {
159            let bytes: SmallVec<[u8; 8]> = [1, 2, 3, 4, 5].iter().cloned().collect();
160            let obj = bytes.clone().into_pyobject(py).unwrap();
161            assert!(obj.is_instance_of::<PyBytes>());
162            let obj = obj.cast_into::<PyBytes>().unwrap();
163            assert_eq!(obj.as_bytes(), &*bytes);
164
165            let nums: SmallVec<[u16; 8]> = [1, 2, 3, 4, 5].iter().cloned().collect();
166            let obj = nums.into_pyobject(py).unwrap();
167            assert!(obj.is_instance_of::<PyList>());
168        });
169    }
170}