pyo3/conversions/std/
slice.rs

1use std::borrow::Cow;
2
3#[cfg(feature = "experimental-inspect")]
4use crate::inspect::types::TypeInfo;
5use crate::{
6    conversion::IntoPyObject,
7    types::{PyByteArray, PyByteArrayMethods, PyBytes},
8    Bound, PyAny, PyErr, PyResult, Python,
9};
10
11impl<'a, 'py, T> IntoPyObject<'py> for &'a [T]
12where
13    &'a T: IntoPyObject<'py>,
14    T: 'a, // MSRV
15{
16    type Target = PyAny;
17    type Output = Bound<'py, Self::Target>;
18    type Error = PyErr;
19
20    /// Turns [`&[u8]`](std::slice) into [`PyBytes`], all other `T`s will be turned into a [`PyList`]
21    ///
22    /// [`PyBytes`]: crate::types::PyBytes
23    /// [`PyList`]: crate::types::PyList
24    #[inline]
25    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
26        <&T>::borrowed_sequence_into_pyobject(self, py, crate::conversion::private::Token)
27    }
28
29    #[cfg(feature = "experimental-inspect")]
30    fn type_output() -> TypeInfo {
31        TypeInfo::union_of(&[
32            TypeInfo::builtin("bytes"),
33            TypeInfo::list_of(<&T>::type_output()),
34        ])
35    }
36}
37
38impl<'a> crate::conversion::FromPyObjectBound<'a, '_> for &'a [u8] {
39    fn from_py_object_bound(obj: crate::Borrowed<'a, '_, PyAny>) -> PyResult<Self> {
40        Ok(obj.downcast::<PyBytes>()?.as_bytes())
41    }
42
43    #[cfg(feature = "experimental-inspect")]
44    fn type_input() -> TypeInfo {
45        TypeInfo::builtin("bytes")
46    }
47}
48
49/// Special-purpose trait impl to efficiently handle both `bytes` and `bytearray`
50///
51/// If the source object is a `bytes` object, the `Cow` will be borrowed and
52/// pointing into the source object, and no copying or heap allocations will happen.
53/// If it is a `bytearray`, its contents will be copied to an owned `Cow`.
54impl<'a> crate::conversion::FromPyObjectBound<'a, '_> for Cow<'a, [u8]> {
55    fn from_py_object_bound(ob: crate::Borrowed<'a, '_, PyAny>) -> PyResult<Self> {
56        if let Ok(bytes) = ob.downcast::<PyBytes>() {
57            return Ok(Cow::Borrowed(bytes.as_bytes()));
58        }
59
60        let byte_array = ob.downcast::<PyByteArray>()?;
61        Ok(Cow::Owned(byte_array.to_vec()))
62    }
63
64    #[cfg(feature = "experimental-inspect")]
65    fn type_input() -> TypeInfo {
66        Self::type_output()
67    }
68}
69
70impl<'py, T> IntoPyObject<'py> for Cow<'_, [T]>
71where
72    T: Clone,
73    for<'a> &'a T: IntoPyObject<'py>,
74{
75    type Target = PyAny;
76    type Output = Bound<'py, Self::Target>;
77    type Error = PyErr;
78
79    /// Turns `Cow<[u8]>` into [`PyBytes`], all other `T`s will be turned into a [`PyList`]
80    ///
81    /// [`PyBytes`]: crate::types::PyBytes
82    /// [`PyList`]: crate::types::PyList
83    #[inline]
84    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
85        <&T>::borrowed_sequence_into_pyobject(self.as_ref(), py, crate::conversion::private::Token)
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use std::borrow::Cow;
92
93    use crate::{
94        conversion::IntoPyObject,
95        ffi,
96        types::{any::PyAnyMethods, PyBytes, PyBytesMethods, PyList},
97        Python,
98    };
99
100    #[test]
101    fn test_extract_bytes() {
102        Python::with_gil(|py| {
103            let py_bytes = py.eval(ffi::c_str!("b'Hello Python'"), None, None).unwrap();
104            let bytes: &[u8] = py_bytes.extract().unwrap();
105            assert_eq!(bytes, b"Hello Python");
106        });
107    }
108
109    #[test]
110    fn test_cow_impl() {
111        Python::with_gil(|py| {
112            let bytes = py.eval(ffi::c_str!(r#"b"foobar""#), None, None).unwrap();
113            let cow = bytes.extract::<Cow<'_, [u8]>>().unwrap();
114            assert_eq!(cow, Cow::<[u8]>::Borrowed(b"foobar"));
115
116            let byte_array = py
117                .eval(ffi::c_str!(r#"bytearray(b"foobar")"#), None, None)
118                .unwrap();
119            let cow = byte_array.extract::<Cow<'_, [u8]>>().unwrap();
120            assert_eq!(cow, Cow::<[u8]>::Owned(b"foobar".to_vec()));
121
122            let something_else_entirely = py.eval(ffi::c_str!("42"), None, None).unwrap();
123            something_else_entirely
124                .extract::<Cow<'_, [u8]>>()
125                .unwrap_err();
126
127            let cow = Cow::<[u8]>::Borrowed(b"foobar").into_pyobject(py).unwrap();
128            assert!(cow.is_instance_of::<PyBytes>());
129
130            let cow = Cow::<[u8]>::Owned(b"foobar".to_vec())
131                .into_pyobject(py)
132                .unwrap();
133            assert!(cow.is_instance_of::<PyBytes>());
134        });
135    }
136
137    #[test]
138    fn test_slice_intopyobject_impl() {
139        Python::with_gil(|py| {
140            let bytes: &[u8] = b"foobar";
141            let obj = bytes.into_pyobject(py).unwrap();
142            assert!(obj.is_instance_of::<PyBytes>());
143            let obj = obj.downcast_into::<PyBytes>().unwrap();
144            assert_eq!(obj.as_bytes(), bytes);
145
146            let nums: &[u16] = &[0, 1, 2, 3];
147            let obj = nums.into_pyobject(py).unwrap();
148            assert!(obj.is_instance_of::<PyList>());
149        });
150    }
151
152    #[test]
153    fn test_cow_intopyobject_impl() {
154        Python::with_gil(|py| {
155            let borrowed_bytes = Cow::<[u8]>::Borrowed(b"foobar");
156            let obj = borrowed_bytes.clone().into_pyobject(py).unwrap();
157            assert!(obj.is_instance_of::<PyBytes>());
158            let obj = obj.downcast_into::<PyBytes>().unwrap();
159            assert_eq!(obj.as_bytes(), &*borrowed_bytes);
160
161            let owned_bytes = Cow::<[u8]>::Owned(b"foobar".to_vec());
162            let obj = owned_bytes.clone().into_pyobject(py).unwrap();
163            assert!(obj.is_instance_of::<PyBytes>());
164            let obj = obj.downcast_into::<PyBytes>().unwrap();
165            assert_eq!(obj.as_bytes(), &*owned_bytes);
166
167            let borrowed_nums = Cow::<[u16]>::Borrowed(&[0, 1, 2, 3]);
168            let obj = borrowed_nums.into_pyobject(py).unwrap();
169            assert!(obj.is_instance_of::<PyList>());
170
171            let owned_nums = Cow::<[u16]>::Owned(vec![0, 1, 2, 3]);
172            let obj = owned_nums.into_pyobject(py).unwrap();
173            assert!(obj.is_instance_of::<PyList>());
174        });
175    }
176}