Skip to main content

pyo3/conversions/std/
slice.rs

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