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    let seq = unsafe {
99        if ffi::PySequence_Check(obj.as_ptr()) != 0 {
100            obj.cast_unchecked::<PySequence>()
101        } else {
102            return Err(CastError::new(obj, PySequence::type_object(obj.py()).into_any()).into());
103        }
104    };
105
106    let mut v = Vec::with_capacity(seq.len().unwrap_or(0));
107    for item in seq.try_iter()? {
108        v.push(item?.extract::<T>().map_err(Into::into)?);
109    }
110    Ok(v)
111}
112
113#[cfg(test)]
114mod tests {
115    use crate::conversion::IntoPyObject;
116    use crate::types::{PyAnyMethods, PyBytes, PyBytesMethods, PyList};
117    use crate::Python;
118
119    #[test]
120    fn test_vec_intopyobject_impl() {
121        Python::attach(|py| {
122            let bytes: Vec<u8> = b"foobar".to_vec();
123            let obj = bytes.clone().into_pyobject(py).unwrap();
124            assert!(obj.is_instance_of::<PyBytes>());
125            let obj = obj.cast_into::<PyBytes>().unwrap();
126            assert_eq!(obj.as_bytes(), &bytes);
127
128            let nums: Vec<u16> = vec![0, 1, 2, 3];
129            let obj = nums.into_pyobject(py).unwrap();
130            assert!(obj.is_instance_of::<PyList>());
131        });
132    }
133
134    #[test]
135    fn test_vec_reference_intopyobject_impl() {
136        Python::attach(|py| {
137            let bytes: Vec<u8> = b"foobar".to_vec();
138            let obj = (&bytes).into_pyobject(py).unwrap();
139            assert!(obj.is_instance_of::<PyBytes>());
140            let obj = obj.cast_into::<PyBytes>().unwrap();
141            assert_eq!(obj.as_bytes(), &bytes);
142
143            let nums: Vec<u16> = vec![0, 1, 2, 3];
144            let obj = (&nums).into_pyobject(py).unwrap();
145            assert!(obj.is_instance_of::<PyList>());
146        });
147    }
148
149    #[test]
150    fn test_strings_cannot_be_extracted_to_vec() {
151        Python::attach(|py| {
152            let v = "London Calling";
153            let ob = v.into_pyobject(py).unwrap();
154
155            assert!(ob.extract::<Vec<String>>().is_err());
156            assert!(ob.extract::<Vec<char>>().is_err());
157        });
158    }
159
160    #[test]
161    fn test_extract_bytes_to_vec() {
162        Python::attach(|py| {
163            let v: Vec<u8> = PyBytes::new(py, b"abc").extract().unwrap();
164            assert_eq!(v, b"abc");
165        });
166    }
167
168    #[test]
169    fn test_extract_tuple_to_vec() {
170        Python::attach(|py| {
171            let v: Vec<i32> = py.eval(c"(1, 2)", None, None).unwrap().extract().unwrap();
172            assert_eq!(v, [1, 2]);
173        });
174    }
175
176    #[test]
177    fn test_extract_range_to_vec() {
178        Python::attach(|py| {
179            let v: Vec<i32> = py
180                .eval(c"range(1, 5)", None, None)
181                .unwrap()
182                .extract()
183                .unwrap();
184            assert_eq!(v, [1, 2, 3, 4]);
185        });
186    }
187
188    #[test]
189    fn test_extract_bytearray_to_vec() {
190        Python::attach(|py| {
191            let v: Vec<u8> = py
192                .eval(c"bytearray(b'abc')", None, None)
193                .unwrap()
194                .extract()
195                .unwrap();
196            assert_eq!(v, b"abc");
197        });
198    }
199}