pyo3/conversions/std/
slice.rs

1use std::borrow::Cow;
2
3#[cfg(feature = "experimental-inspect")]
4use crate::inspect::types::TypeInfo;
5use crate::{
6    types::{PyByteArray, PyByteArrayMethods, PyBytes},
7    IntoPy, Py, PyAny, PyObject, PyResult, Python, ToPyObject,
8};
9
10impl<'a> IntoPy<PyObject> for &'a [u8] {
11    fn into_py(self, py: Python<'_>) -> PyObject {
12        PyBytes::new_bound(py, self).unbind().into()
13    }
14
15    #[cfg(feature = "experimental-inspect")]
16    fn type_output() -> TypeInfo {
17        TypeInfo::builtin("bytes")
18    }
19}
20
21#[cfg(feature = "gil-refs")]
22impl<'py> crate::FromPyObject<'py> for &'py [u8] {
23    fn extract_bound(obj: &crate::Bound<'py, PyAny>) -> PyResult<Self> {
24        Ok(obj.clone().into_gil_ref().downcast::<PyBytes>()?.as_bytes())
25    }
26
27    #[cfg(feature = "experimental-inspect")]
28    fn type_input() -> TypeInfo {
29        Self::type_output()
30    }
31}
32
33#[cfg(not(feature = "gil-refs"))]
34impl<'a> crate::conversion::FromPyObjectBound<'a, '_> for &'a [u8] {
35    fn from_py_object_bound(obj: crate::Borrowed<'a, '_, PyAny>) -> PyResult<Self> {
36        Ok(obj.downcast::<PyBytes>()?.as_bytes())
37    }
38
39    #[cfg(feature = "experimental-inspect")]
40    fn type_input() -> TypeInfo {
41        Self::type_output()
42    }
43}
44
45/// Special-purpose trait impl to efficiently handle both `bytes` and `bytearray`
46///
47/// If the source object is a `bytes` object, the `Cow` will be borrowed and
48/// pointing into the source object, and no copying or heap allocations will happen.
49/// If it is a `bytearray`, its contents will be copied to an owned `Cow`.
50#[cfg(feature = "gil-refs")]
51impl<'py> crate::FromPyObject<'py> for Cow<'py, [u8]> {
52    fn extract_bound(ob: &crate::Bound<'py, PyAny>) -> PyResult<Self> {
53        use crate::types::PyAnyMethods;
54        if let Ok(bytes) = ob.downcast::<PyBytes>() {
55            return Ok(Cow::Borrowed(bytes.clone().into_gil_ref().as_bytes()));
56        }
57
58        let byte_array = ob.downcast::<PyByteArray>()?;
59        Ok(Cow::Owned(byte_array.to_vec()))
60    }
61}
62
63#[cfg(not(feature = "gil-refs"))]
64impl<'a> crate::conversion::FromPyObjectBound<'a, '_> for Cow<'a, [u8]> {
65    fn from_py_object_bound(ob: crate::Borrowed<'a, '_, PyAny>) -> PyResult<Self> {
66        if let Ok(bytes) = ob.downcast::<PyBytes>() {
67            return Ok(Cow::Borrowed(bytes.as_bytes()));
68        }
69
70        let byte_array = ob.downcast::<PyByteArray>()?;
71        Ok(Cow::Owned(byte_array.to_vec()))
72    }
73
74    #[cfg(feature = "experimental-inspect")]
75    fn type_input() -> TypeInfo {
76        Self::type_output()
77    }
78}
79
80impl ToPyObject for Cow<'_, [u8]> {
81    fn to_object(&self, py: Python<'_>) -> Py<PyAny> {
82        PyBytes::new_bound(py, self.as_ref()).into()
83    }
84}
85
86impl IntoPy<Py<PyAny>> for Cow<'_, [u8]> {
87    fn into_py(self, py: Python<'_>) -> Py<PyAny> {
88        self.to_object(py)
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use std::borrow::Cow;
95
96    use crate::{
97        types::{any::PyAnyMethods, PyBytes},
98        Python, ToPyObject,
99    };
100
101    #[test]
102    fn test_extract_bytes() {
103        Python::with_gil(|py| {
104            let py_bytes = py.eval_bound("b'Hello Python'", None, None).unwrap();
105            let bytes: &[u8] = py_bytes.extract().unwrap();
106            assert_eq!(bytes, b"Hello Python");
107        });
108    }
109
110    #[test]
111    fn test_cow_impl() {
112        Python::with_gil(|py| {
113            let bytes = py.eval_bound(r#"b"foobar""#, None, None).unwrap();
114            let cow = bytes.extract::<Cow<'_, [u8]>>().unwrap();
115            assert_eq!(cow, Cow::<[u8]>::Borrowed(b"foobar"));
116
117            let byte_array = py
118                .eval_bound(r#"bytearray(b"foobar")"#, None, None)
119                .unwrap();
120            let cow = byte_array.extract::<Cow<'_, [u8]>>().unwrap();
121            assert_eq!(cow, Cow::<[u8]>::Owned(b"foobar".to_vec()));
122
123            let something_else_entirely = py.eval_bound("42", None, None).unwrap();
124            something_else_entirely
125                .extract::<Cow<'_, [u8]>>()
126                .unwrap_err();
127
128            let cow = Cow::<[u8]>::Borrowed(b"foobar").to_object(py);
129            assert!(cow.bind(py).is_instance_of::<PyBytes>());
130
131            let cow = Cow::<[u8]>::Owned(b"foobar".to_vec()).to_object(py);
132            assert!(cow.bind(py).is_instance_of::<PyBytes>());
133        });
134    }
135}