Skip to main content

pyo3/conversions/std/
vec.rs

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