Skip to main content

pyo3/conversions/std/
slice.rs

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