1use crate::{ffi, FromPyObject, IntoPy, Py, PyAny, PyResult, Python, ToPyObject};
2use std::borrow::Cow;
3use std::ops::Index;
4use std::os::raw::c_char;
5use std::slice::SliceIndex;
6use std::str;
7
8use super::bytearray::PyByteArray;
9
10/// Represents a Python `bytes` object.
11///
12/// This type is immutable.
13#[repr(transparent)]
14pub struct PyBytes(PyAny);
15
16pyobject_native_type_core!(PyBytes, pyobject_native_static_type_object!(ffi::PyBytes_Type), #checkfunction=ffi::PyBytes_Check);
17
18impl PyBytes {
19 /// Creates a new Python bytestring object.
20 /// The bytestring is initialized by copying the data from the `&[u8]`.
21 ///
22 /// Panics if out of memory.
23 pub fn new<'p>(py: Python<'p>, s: &[u8]) -> &'p PyBytes {
24 let ptr = s.as_ptr() as *const c_char;
25 let len = s.len() as ffi::Py_ssize_t;
26 unsafe { py.from_owned_ptr(ffi::PyBytes_FromStringAndSize(ptr, len)) }
27 }
28
29 /// Creates a new Python `bytes` object with an `init` closure to write its contents.
30 /// Before calling `init` the bytes' contents are zero-initialised.
31 /// * If Python raises a MemoryError on the allocation, `new_with` will return
32 /// it inside `Err`.
33 /// * If `init` returns `Err(e)`, `new_with` will return `Err(e)`.
34 /// * If `init` returns `Ok(())`, `new_with` will return `Ok(&PyBytes)`.
35 ///
36 /// # Examples
37 ///
38 /// ```
39 /// use pyo3::{prelude::*, types::PyBytes};
40 ///
41 /// # fn main() -> PyResult<()> {
42 /// Python::with_gil(|py| -> PyResult<()> {
43 /// let py_bytes = PyBytes::new_with(py, 10, |bytes: &mut [u8]| {
44 /// bytes.copy_from_slice(b"Hello Rust");
45 /// Ok(())
46 /// })?;
47 /// let bytes: &[u8] = FromPyObject::extract(py_bytes)?;
48 /// assert_eq!(bytes, b"Hello Rust");
49 /// Ok(())
50 /// })
51 /// # }
52 /// ```
53 pub fn new_with<F>(py: Python<'_>, len: usize, init: F) -> PyResult<&PyBytes>
54 where
55 F: FnOnce(&mut [u8]) -> PyResult<()>,
56 {
57 unsafe {
58 let pyptr = ffi::PyBytes_FromStringAndSize(std::ptr::null(), len as ffi::Py_ssize_t);
59 // Check for an allocation error and return it
60 let pypybytes: Py<PyBytes> = Py::from_owned_ptr_or_err(py, pyptr)?;
61 let buffer: *mut u8 = ffi::PyBytes_AsString(pyptr).cast();
62 debug_assert!(!buffer.is_null());
63 // Zero-initialise the uninitialised bytestring
64 std::ptr::write_bytes(buffer, 0u8, len);
65 // (Further) Initialise the bytestring in init
66 // If init returns an Err, pypybytearray will automatically deallocate the buffer
67 init(std::slice::from_raw_parts_mut(buffer, len)).map(|_| pypybytes.into_ref(py))
68 }
69 }
70
71 /// Creates a new Python byte string object from a raw pointer and length.
72 ///
73 /// Panics if out of memory.
74 ///
75 /// # Safety
76 ///
77 /// This function dereferences the raw pointer `ptr` as the
78 /// leading pointer of a slice of length `len`. [As with
79 /// `std::slice::from_raw_parts`, this is
80 /// unsafe](https://doc.rust-lang.org/std/slice/fn.from_raw_parts.html#safety).
81 pub unsafe fn from_ptr(py: Python<'_>, ptr: *const u8, len: usize) -> &PyBytes {
82 py.from_owned_ptr(ffi::PyBytes_FromStringAndSize(
83 ptr as *const _,
84 len as isize,
85 ))
86 }
87
88 /// Gets the Python string as a byte slice.
89 #[inline]
90 pub fn as_bytes(&self) -> &[u8] {
91 unsafe {
92 let buffer = ffi::PyBytes_AsString(self.as_ptr()) as *const u8;
93 let length = ffi::PyBytes_Size(self.as_ptr()) as usize;
94 debug_assert!(!buffer.is_null());
95 std::slice::from_raw_parts(buffer, length)
96 }
97 }
98}
99
100impl Py<PyBytes> {
101 /// Gets the Python bytes as a byte slice. Because Python bytes are
102 /// immutable, the result may be used for as long as the reference to
103 /// `self` is held, including when the GIL is released.
104 pub fn as_bytes<'a>(&'a self, _py: Python<'_>) -> &'a [u8] {
105 // py is required here because `PyBytes_AsString` and `PyBytes_Size`
106 // can both technically raise exceptions which require the GIL to be
107 // held. The only circumstance in which they raise is if the value
108 // isn't really a `PyBytes`, but better safe than sorry.
109 unsafe {
110 let buffer: *const u8 = ffi::PyBytes_AsString(self.as_ptr()) as *const u8;
111 let length: usize = ffi::PyBytes_Size(self.as_ptr()) as usize;
112 debug_assert!(!buffer.is_null());
113 std::slice::from_raw_parts(data:buffer, len:length)
114 }
115 }
116}
117
118/// This is the same way [Vec] is indexed.
119impl<I: SliceIndex<[u8]>> Index<I> for PyBytes {
120 type Output = I::Output;
121
122 fn index(&self, index: I) -> &Self::Output {
123 &self.as_bytes()[index]
124 }
125}
126
127/// Special-purpose trait impl to efficiently handle both `bytes` and `bytearray`
128///
129/// If the source object is a `bytes` object, the `Cow` will be borrowed and
130/// pointing into the source object, and no copying or heap allocations will happen.
131/// If it is a `bytearray`, its contents will be copied to an owned `Cow`.
132impl<'source> FromPyObject<'source> for Cow<'source, [u8]> {
133 fn extract(ob: &'source PyAny) -> PyResult<Self> {
134 if let Ok(bytes: &PyBytes) = ob.downcast::<PyBytes>() {
135 return Ok(Cow::Borrowed(bytes.as_bytes()));
136 }
137
138 let byte_array: &PyByteArray = ob.downcast::<PyByteArray>()?;
139 Ok(Cow::Owned(byte_array.to_vec()))
140 }
141}
142
143impl ToPyObject for Cow<'_, [u8]> {
144 fn to_object(&self, py: Python<'_>) -> Py<PyAny> {
145 PyBytes::new(py, self.as_ref()).into()
146 }
147}
148
149impl IntoPy<Py<PyAny>> for Cow<'_, [u8]> {
150 fn into_py(self, py: Python<'_>) -> Py<PyAny> {
151 self.to_object(py)
152 }
153}
154
155#[cfg(test)]
156mod tests {
157 use super::*;
158
159 #[test]
160 fn test_bytes_index() {
161 Python::with_gil(|py| {
162 let bytes = PyBytes::new(py, b"Hello World");
163 assert_eq!(bytes[1], b'e');
164 });
165 }
166
167 #[test]
168 fn test_bytes_new_with() -> super::PyResult<()> {
169 Python::with_gil(|py| -> super::PyResult<()> {
170 let py_bytes = PyBytes::new_with(py, 10, |b: &mut [u8]| {
171 b.copy_from_slice(b"Hello Rust");
172 Ok(())
173 })?;
174 let bytes: &[u8] = FromPyObject::extract(py_bytes)?;
175 assert_eq!(bytes, b"Hello Rust");
176 Ok(())
177 })
178 }
179
180 #[test]
181 fn test_bytes_new_with_zero_initialised() -> super::PyResult<()> {
182 Python::with_gil(|py| -> super::PyResult<()> {
183 let py_bytes = PyBytes::new_with(py, 10, |_b: &mut [u8]| Ok(()))?;
184 let bytes: &[u8] = FromPyObject::extract(py_bytes)?;
185 assert_eq!(bytes, &[0; 10]);
186 Ok(())
187 })
188 }
189
190 #[test]
191 fn test_bytes_new_with_error() {
192 use crate::exceptions::PyValueError;
193 Python::with_gil(|py| {
194 let py_bytes_result = PyBytes::new_with(py, 10, |_b: &mut [u8]| {
195 Err(PyValueError::new_err("Hello Crustaceans!"))
196 });
197 assert!(py_bytes_result.is_err());
198 assert!(py_bytes_result
199 .err()
200 .unwrap()
201 .is_instance_of::<PyValueError>(py));
202 });
203 }
204
205 #[test]
206 fn test_cow_impl() {
207 Python::with_gil(|py| {
208 let bytes = py.eval(r#"b"foobar""#, None, None).unwrap();
209 let cow = bytes.extract::<Cow<'_, [u8]>>().unwrap();
210 assert_eq!(cow, Cow::<[u8]>::Borrowed(b"foobar"));
211
212 let byte_array = py.eval(r#"bytearray(b"foobar")"#, None, None).unwrap();
213 let cow = byte_array.extract::<Cow<'_, [u8]>>().unwrap();
214 assert_eq!(cow, Cow::<[u8]>::Owned(b"foobar".to_vec()));
215
216 let something_else_entirely = py.eval("42", None, None).unwrap();
217 something_else_entirely
218 .extract::<Cow<'_, [u8]>>()
219 .unwrap_err();
220
221 let cow = Cow::<[u8]>::Borrowed(b"foobar").to_object(py);
222 assert!(cow.as_ref(py).is_instance_of::<PyBytes>());
223
224 let cow = Cow::<[u8]>::Owned(b"foobar".to_vec()).to_object(py);
225 assert!(cow.as_ref(py).is_instance_of::<PyBytes>());
226 });
227 }
228}
229