Skip to main content

pyo3/conversions/std/
vec.rs

1#[cfg(feature = "experimental-inspect")]
2use crate::inspect::types::TypeInfo;
3#[cfg(feature = "experimental-inspect")]
4use crate::inspect::{type_hint_subscript, PyStaticExpr};
5use crate::{
6    conversion::{FromPyObject, FromPyObjectOwned, FromPyObjectSequence, IntoPyObject},
7    exceptions::PyTypeError,
8    ffi,
9    types::{PyAnyMethods, PySequence, PyString},
10    Borrowed, CastError, PyResult, PyTypeInfo,
11};
12use crate::{Bound, PyAny, PyErr, Python};
13
14impl<'py, T> IntoPyObject<'py> for Vec<T>
15where
16    T: IntoPyObject<'py>,
17{
18    type Target = PyAny;
19    type Output = Bound<'py, Self::Target>;
20    type Error = PyErr;
21
22    #[cfg(feature = "experimental-inspect")]
23    const OUTPUT_TYPE: PyStaticExpr = T::SEQUENCE_OUTPUT_TYPE;
24
25    /// Turns [`Vec<u8>`] into [`PyBytes`], all other `T`s will be turned into a [`PyList`]
26    ///
27    /// [`PyBytes`]: crate::types::PyBytes
28    /// [`PyList`]: crate::types::PyList
29    #[inline]
30    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
31        T::owned_sequence_into_pyobject(self, py, crate::conversion::private::Token)
32    }
33
34    #[cfg(feature = "experimental-inspect")]
35    fn type_output() -> TypeInfo {
36        TypeInfo::list_of(T::type_output())
37    }
38}
39
40impl<'a, 'py, T> IntoPyObject<'py> for &'a Vec<T>
41where
42    &'a T: IntoPyObject<'py>,
43{
44    type Target = PyAny;
45    type Output = Bound<'py, Self::Target>;
46    type Error = PyErr;
47
48    #[cfg(feature = "experimental-inspect")]
49    const OUTPUT_TYPE: PyStaticExpr = <&[T]>::OUTPUT_TYPE;
50
51    #[inline]
52    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
53        // NB: we could actually not cast to `PyAny`, which would be nice for
54        // `&Vec<u8>`, but that'd be inconsistent with the `IntoPyObject` impl
55        // above which always returns a `PyAny` for `Vec<T>`.
56        self.as_slice().into_pyobject(py).map(Bound::into_any)
57    }
58
59    #[cfg(feature = "experimental-inspect")]
60    fn type_output() -> TypeInfo {
61        TypeInfo::list_of(<&T>::type_output())
62    }
63}
64
65impl<'py, T> FromPyObject<'_, 'py> for Vec<T>
66where
67    T: FromPyObjectOwned<'py>,
68{
69    type Error = PyErr;
70
71    #[cfg(feature = "experimental-inspect")]
72    const INPUT_TYPE: PyStaticExpr = type_hint_subscript!(PySequence::TYPE_HINT, T::INPUT_TYPE);
73
74    fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {
75        if let Some(extractor) = T::sequence_extractor(obj, crate::conversion::private::Token) {
76            return Ok(extractor.to_vec());
77        }
78
79        if obj.is_instance_of::<PyString>() {
80            return Err(PyTypeError::new_err("Can't extract `str` to `Vec`"));
81        }
82
83        extract_sequence(obj)
84    }
85
86    #[cfg(feature = "experimental-inspect")]
87    fn type_input() -> TypeInfo {
88        TypeInfo::sequence_of(T::type_input())
89    }
90}
91
92fn extract_sequence<'py, T>(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<Vec<T>>
93where
94    T: FromPyObjectOwned<'py>,
95{
96    // Types that pass `PySequence_Check` usually implement enough of the sequence protocol
97    // to support this function and if not, we will only fail extraction safely.
98    if unsafe { ffi::PySequence_Check(obj.as_ptr()) } == 0 {
99        return Err(CastError::new(obj, PySequence::type_object(obj.py()).into_any()).into());
100    }
101
102    let mut v = Vec::with_capacity(obj.len().unwrap_or(0));
103    for item in obj.try_iter()? {
104        v.push(item?.extract::<T>().map_err(Into::into)?);
105    }
106    Ok(v)
107}
108
109#[cfg(test)]
110mod tests {
111    use crate::conversion::IntoPyObject;
112    use crate::types::{PyAnyMethods, PyBytes, PyBytesMethods, PyList};
113    use crate::Python;
114
115    #[test]
116    fn test_vec_intopyobject_impl() {
117        Python::attach(|py| {
118            let bytes: Vec<u8> = b"foobar".to_vec();
119            let obj = bytes.clone().into_pyobject(py).unwrap();
120            assert!(obj.is_instance_of::<PyBytes>());
121            let obj = obj.cast_into::<PyBytes>().unwrap();
122            assert_eq!(obj.as_bytes(), &bytes);
123
124            let nums: Vec<u16> = vec![0, 1, 2, 3];
125            let obj = nums.into_pyobject(py).unwrap();
126            assert!(obj.is_instance_of::<PyList>());
127        });
128    }
129
130    #[test]
131    fn test_vec_reference_intopyobject_impl() {
132        Python::attach(|py| {
133            let bytes: Vec<u8> = b"foobar".to_vec();
134            let obj = (&bytes).into_pyobject(py).unwrap();
135            assert!(obj.is_instance_of::<PyBytes>());
136            let obj = obj.cast_into::<PyBytes>().unwrap();
137            assert_eq!(obj.as_bytes(), &bytes);
138
139            let nums: Vec<u16> = vec![0, 1, 2, 3];
140            let obj = (&nums).into_pyobject(py).unwrap();
141            assert!(obj.is_instance_of::<PyList>());
142        });
143    }
144
145    #[test]
146    fn test_strings_cannot_be_extracted_to_vec() {
147        Python::attach(|py| {
148            let v = "London Calling";
149            let ob = v.into_pyobject(py).unwrap();
150
151            assert!(ob.extract::<Vec<String>>().is_err());
152            assert!(ob.extract::<Vec<char>>().is_err());
153        });
154    }
155
156    #[test]
157    fn test_extract_bytes_to_vec() {
158        Python::attach(|py| {
159            let v: Vec<u8> = PyBytes::new(py, b"abc").extract().unwrap();
160            assert_eq!(v, b"abc");
161        });
162    }
163
164    #[test]
165    fn test_extract_tuple_to_vec() {
166        Python::attach(|py| {
167            let v: Vec<i32> = py.eval(c"(1, 2)", None, None).unwrap().extract().unwrap();
168            assert_eq!(v, [1, 2]);
169        });
170    }
171
172    #[test]
173    fn test_extract_range_to_vec() {
174        Python::attach(|py| {
175            let v: Vec<i32> = py
176                .eval(c"range(1, 5)", None, None)
177                .unwrap()
178                .extract()
179                .unwrap();
180            assert_eq!(v, [1, 2, 3, 4]);
181        });
182    }
183
184    #[test]
185    fn test_extract_bytearray_to_vec() {
186        Python::attach(|py| {
187            let v: Vec<u8> = py
188                .eval(c"bytearray(b'abc')", None, None)
189                .unwrap()
190                .extract()
191                .unwrap();
192            assert_eq!(v, b"abc");
193        });
194    }
195}