1use crate::err::{PyErr, PyResult};
2use crate::ffi_ptr_ext::FfiPtrExt;
3use crate::instance::{Borrowed, Bound};
4use crate::py_result_ext::PyResultExt;
5use crate::types::any::PyAnyMethods;
6use crate::{ffi, PyAny, Python};
7use std::slice;
8
9/// Represents a Python `bytearray`.
10///
11/// Values of this type are accessed via PyO3's smart pointers, e.g. as
12/// [`Py<PyByteArray>`][crate::Py] or [`Bound<'py, PyByteArray>`][Bound].
13///
14/// For APIs available on `bytearray` objects, see the [`PyByteArrayMethods`] trait which is implemented for
15/// [`Bound<'py, PyByteArray>`][Bound].
16#[repr(transparent)]
17pub struct PyByteArray(PyAny);
18
19pyobject_native_type_core!(PyByteArray, pyobject_native_static_type_object!(ffi::PyByteArray_Type), #checkfunction=ffi::PyByteArray_Check);
20
21impl PyByteArray {
22 /// Creates a new Python bytearray object.
23 ///
24 /// The byte string is initialized by copying the data from the `&[u8]`.
25 pub fn new<'py>(py: Python<'py>, src: &[u8]) -> Bound<'py, PyByteArray> {
26 let ptr = src.as_ptr().cast();
27 let len = src.len() as ffi::Py_ssize_t;
28 unsafe {
29 ffi::PyByteArray_FromStringAndSize(ptr, len)
30 .assume_owned(py)
31 .downcast_into_unchecked()
32 }
33 }
34
35 /// Deprecated name for [`PyByteArray::new`].
36 #[deprecated(since = "0.23.0", note = "renamed to `PyByteArray::new`")]
37 #[inline]
38 pub fn new_bound<'py>(py: Python<'py>, src: &[u8]) -> Bound<'py, PyByteArray> {
39 Self::new(py, src)
40 }
41
42 /// Creates a new Python `bytearray` object with an `init` closure to write its contents.
43 /// Before calling `init` the bytearray is zero-initialised.
44 /// * If Python raises a MemoryError on the allocation, `new_with` will return
45 /// it inside `Err`.
46 /// * If `init` returns `Err(e)`, `new_with` will return `Err(e)`.
47 /// * If `init` returns `Ok(())`, `new_with` will return `Ok(&PyByteArray)`.
48 ///
49 /// # Examples
50 ///
51 /// ```
52 /// use pyo3::{prelude::*, types::PyByteArray};
53 ///
54 /// # fn main() -> PyResult<()> {
55 /// Python::with_gil(|py| -> PyResult<()> {
56 /// let py_bytearray = PyByteArray::new_with(py, 10, |bytes: &mut [u8]| {
57 /// bytes.copy_from_slice(b"Hello Rust");
58 /// Ok(())
59 /// })?;
60 /// let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
61 /// assert_eq!(bytearray, b"Hello Rust");
62 /// Ok(())
63 /// })
64 /// # }
65 /// ```
66 pub fn new_with<F>(py: Python<'_>, len: usize, init: F) -> PyResult<Bound<'_, PyByteArray>>
67 where
68 F: FnOnce(&mut [u8]) -> PyResult<()>,
69 {
70 unsafe {
71 // Allocate buffer and check for an error
72 let pybytearray: Bound<'_, Self> =
73 ffi::PyByteArray_FromStringAndSize(std::ptr::null(), len as ffi::Py_ssize_t)
74 .assume_owned_or_err(py)?
75 .downcast_into_unchecked();
76
77 let buffer: *mut u8 = ffi::PyByteArray_AsString(pybytearray.as_ptr()).cast();
78 debug_assert!(!buffer.is_null());
79 // Zero-initialise the uninitialised bytearray
80 std::ptr::write_bytes(buffer, 0u8, len);
81 // (Further) Initialise the bytearray in init
82 // If init returns an Err, pypybytearray will automatically deallocate the buffer
83 init(std::slice::from_raw_parts_mut(buffer, len)).map(|_| pybytearray)
84 }
85 }
86
87 /// Deprecated name for [`PyByteArray::new_with`].
88 #[deprecated(since = "0.23.0", note = "renamed to `PyByteArray::new_with`")]
89 #[inline]
90 pub fn new_bound_with<F>(
91 py: Python<'_>,
92 len: usize,
93 init: F,
94 ) -> PyResult<Bound<'_, PyByteArray>>
95 where
96 F: FnOnce(&mut [u8]) -> PyResult<()>,
97 {
98 Self::new_with(py, len, init)
99 }
100
101 /// Creates a new Python `bytearray` object from another Python object that
102 /// implements the buffer protocol.
103 pub fn from<'py>(src: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyByteArray>> {
104 unsafe {
105 ffi::PyByteArray_FromObject(src.as_ptr())
106 .assume_owned_or_err(src.py())
107 .downcast_into_unchecked()
108 }
109 }
110
111 ///Deprecated name for [`PyByteArray::from`].
112 #[deprecated(since = "0.23.0", note = "renamed to `PyByteArray::from`")]
113 #[inline]
114 pub fn from_bound<'py>(src: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyByteArray>> {
115 Self::from(src)
116 }
117}
118
119/// Implementation of functionality for [`PyByteArray`].
120///
121/// These methods are defined for the `Bound<'py, PyByteArray>` smart pointer, so to use method call
122/// syntax these methods are separated into a trait, because stable Rust does not yet support
123/// `arbitrary_self_types`.
124#[doc(alias = "PyByteArray")]
125pub trait PyByteArrayMethods<'py>: crate::sealed::Sealed {
126 /// Gets the length of the bytearray.
127 fn len(&self) -> usize;
128
129 /// Checks if the bytearray is empty.
130 fn is_empty(&self) -> bool;
131
132 /// Gets the start of the buffer containing the contents of the bytearray.
133 ///
134 /// # Safety
135 ///
136 /// See the safety requirements of [`PyByteArrayMethods::as_bytes`] and [`PyByteArrayMethods::as_bytes_mut`].
137 fn data(&self) -> *mut u8;
138
139 /// Extracts a slice of the `ByteArray`'s entire buffer.
140 ///
141 /// # Safety
142 ///
143 /// Mutation of the `bytearray` invalidates the slice. If it is used afterwards, the behavior is
144 /// undefined.
145 ///
146 /// These mutations may occur in Python code as well as from Rust:
147 /// - Calling methods like [`PyByteArrayMethods::as_bytes_mut`] and [`PyByteArrayMethods::resize`] will
148 /// invalidate the slice.
149 /// - Actions like dropping objects or raising exceptions can invoke `__del__`methods or signal
150 /// handlers, which may execute arbitrary Python code. This means that if Python code has a
151 /// reference to the `bytearray` you cannot safely use the vast majority of PyO3's API whilst
152 /// using the slice.
153 ///
154 /// As a result, this slice should only be used for short-lived operations without executing any
155 /// Python code, such as copying into a Vec.
156 ///
157 /// # Examples
158 ///
159 /// ```rust
160 /// use pyo3::prelude::*;
161 /// use pyo3::exceptions::PyRuntimeError;
162 /// use pyo3::types::PyByteArray;
163 ///
164 /// #[pyfunction]
165 /// fn a_valid_function(bytes: &Bound<'_, PyByteArray>) -> PyResult<()> {
166 /// let section = {
167 /// // SAFETY: We promise to not let the interpreter regain control
168 /// // or invoke any PyO3 APIs while using the slice.
169 /// let slice = unsafe { bytes.as_bytes() };
170 ///
171 /// // Copy only a section of `bytes` while avoiding
172 /// // `to_vec` which copies the entire thing.
173 /// let section = slice
174 /// .get(6..11)
175 /// .ok_or_else(|| PyRuntimeError::new_err("input is not long enough"))?;
176 /// Vec::from(section)
177 /// };
178 ///
179 /// // Now we can do things with `section` and call PyO3 APIs again.
180 /// // ...
181 /// # assert_eq!(&section, b"world");
182 ///
183 /// Ok(())
184 /// }
185 /// # fn main() -> PyResult<()> {
186 /// # Python::with_gil(|py| -> PyResult<()> {
187 /// # let fun = wrap_pyfunction!(a_valid_function, py)?;
188 /// # let locals = pyo3::types::PyDict::new(py);
189 /// # locals.set_item("a_valid_function", fun)?;
190 /// #
191 /// # py.run(pyo3::ffi::c_str!(
192 /// # r#"b = bytearray(b"hello world")
193 /// # a_valid_function(b)
194 /// #
195 /// # try:
196 /// # a_valid_function(bytearray())
197 /// # except RuntimeError as e:
198 /// # assert str(e) == 'input is not long enough'"#),
199 /// # None,
200 /// # Some(&locals),
201 /// # )?;
202 /// #
203 /// # Ok(())
204 /// # })
205 /// # }
206 /// ```
207 ///
208 /// # Incorrect usage
209 ///
210 /// The following `bug` function is unsound ⚠️
211 ///
212 /// ```rust,no_run
213 /// # use pyo3::prelude::*;
214 /// # use pyo3::types::PyByteArray;
215 ///
216 /// # #[allow(dead_code)]
217 /// #[pyfunction]
218 /// fn bug(py: Python<'_>, bytes: &Bound<'_, PyByteArray>) {
219 /// let slice = unsafe { bytes.as_bytes() };
220 ///
221 /// // This explicitly yields control back to the Python interpreter...
222 /// // ...but it's not always this obvious. Many things do this implicitly.
223 /// py.allow_threads(|| {
224 /// // Python code could be mutating through its handle to `bytes`,
225 /// // which makes reading it a data race, which is undefined behavior.
226 /// println!("{:?}", slice[0]);
227 /// });
228 ///
229 /// // Python code might have mutated it, so we can not rely on the slice
230 /// // remaining valid. As such this is also undefined behavior.
231 /// println!("{:?}", slice[0]);
232 /// }
233 /// ```
234 unsafe fn as_bytes(&self) -> &[u8];
235
236 /// Extracts a mutable slice of the `ByteArray`'s entire buffer.
237 ///
238 /// # Safety
239 ///
240 /// Any other accesses of the `bytearray`'s buffer invalidate the slice. If it is used
241 /// afterwards, the behavior is undefined. The safety requirements of [`PyByteArrayMethods::as_bytes`]
242 /// apply to this function as well.
243 #[allow(clippy::mut_from_ref)]
244 unsafe fn as_bytes_mut(&self) -> &mut [u8];
245
246 /// Copies the contents of the bytearray to a Rust vector.
247 ///
248 /// # Examples
249 ///
250 /// ```
251 /// # use pyo3::prelude::*;
252 /// # use pyo3::types::PyByteArray;
253 /// # Python::with_gil(|py| {
254 /// let bytearray = PyByteArray::new(py, b"Hello World.");
255 /// let mut copied_message = bytearray.to_vec();
256 /// assert_eq!(b"Hello World.", copied_message.as_slice());
257 ///
258 /// copied_message[11] = b'!';
259 /// assert_eq!(b"Hello World!", copied_message.as_slice());
260 ///
261 /// pyo3::py_run!(py, bytearray, "assert bytearray == b'Hello World.'");
262 /// # });
263 /// ```
264 fn to_vec(&self) -> Vec<u8>;
265
266 /// Resizes the bytearray object to the new length `len`.
267 ///
268 /// Note that this will invalidate any pointers obtained by [PyByteArrayMethods::data], as well as
269 /// any (unsafe) slices obtained from [PyByteArrayMethods::as_bytes] and [PyByteArrayMethods::as_bytes_mut].
270 fn resize(&self, len: usize) -> PyResult<()>;
271}
272
273impl<'py> PyByteArrayMethods<'py> for Bound<'py, PyByteArray> {
274 #[inline]
275 fn len(&self) -> usize {
276 // non-negative Py_ssize_t should always fit into Rust usize
277 unsafe { ffi::PyByteArray_Size(self.as_ptr()) as usize }
278 }
279
280 fn is_empty(&self) -> bool {
281 self.len() == 0
282 }
283
284 fn data(&self) -> *mut u8 {
285 self.as_borrowed().data()
286 }
287
288 unsafe fn as_bytes(&self) -> &[u8] {
289 unsafe { self.as_borrowed().as_bytes() }
290 }
291
292 #[allow(clippy::mut_from_ref)]
293 unsafe fn as_bytes_mut(&self) -> &mut [u8] {
294 unsafe { self.as_borrowed().as_bytes_mut() }
295 }
296
297 fn to_vec(&self) -> Vec<u8> {
298 unsafe { self.as_bytes() }.to_vec()
299 }
300
301 fn resize(&self, len: usize) -> PyResult<()> {
302 unsafe {
303 let result = ffi::PyByteArray_Resize(self.as_ptr(), len as ffi::Py_ssize_t);
304 if result == 0 {
305 Ok(())
306 } else {
307 Err(PyErr::fetch(self.py()))
308 }
309 }
310 }
311}
312
313impl<'a> Borrowed<'a, '_, PyByteArray> {
314 fn data(&self) -> *mut u8 {
315 unsafe { ffi::PyByteArray_AsString(self.as_ptr()).cast() }
316 }
317
318 #[allow(clippy::wrong_self_convention)]
319 unsafe fn as_bytes(self) -> &'a [u8] {
320 unsafe { slice::from_raw_parts(self.data(), self.len()) }
321 }
322
323 #[allow(clippy::wrong_self_convention)]
324 unsafe fn as_bytes_mut(self) -> &'a mut [u8] {
325 unsafe { slice::from_raw_parts_mut(self.data(), self.len()) }
326 }
327}
328
329impl<'py> TryFrom<&Bound<'py, PyAny>> for Bound<'py, PyByteArray> {
330 type Error = crate::PyErr;
331
332 /// Creates a new Python `bytearray` object from another Python object that
333 /// implements the buffer protocol.
334 fn try_from(value: &Bound<'py, PyAny>) -> Result<Self, Self::Error> {
335 PyByteArray::from(src:value)
336 }
337}
338
339#[cfg(test)]
340mod tests {
341 use crate::types::{PyAnyMethods, PyByteArray, PyByteArrayMethods};
342 use crate::{exceptions, Bound, PyAny, PyObject, Python};
343
344 #[test]
345 fn test_len() {
346 Python::with_gil(|py| {
347 let src = b"Hello Python";
348 let bytearray = PyByteArray::new(py, src);
349 assert_eq!(src.len(), bytearray.len());
350 });
351 }
352
353 #[test]
354 fn test_as_bytes() {
355 Python::with_gil(|py| {
356 let src = b"Hello Python";
357 let bytearray = PyByteArray::new(py, src);
358
359 let slice = unsafe { bytearray.as_bytes() };
360 assert_eq!(src, slice);
361 assert_eq!(bytearray.data() as *const _, slice.as_ptr());
362 });
363 }
364
365 #[test]
366 fn test_as_bytes_mut() {
367 Python::with_gil(|py| {
368 let src = b"Hello Python";
369 let bytearray = PyByteArray::new(py, src);
370
371 let slice = unsafe { bytearray.as_bytes_mut() };
372 assert_eq!(src, slice);
373 assert_eq!(bytearray.data(), slice.as_mut_ptr());
374
375 slice[0..5].copy_from_slice(b"Hi...");
376
377 assert_eq!(bytearray.str().unwrap(), "bytearray(b'Hi... Python')");
378 });
379 }
380
381 #[test]
382 fn test_to_vec() {
383 Python::with_gil(|py| {
384 let src = b"Hello Python";
385 let bytearray = PyByteArray::new(py, src);
386
387 let vec = bytearray.to_vec();
388 assert_eq!(src, vec.as_slice());
389 });
390 }
391
392 #[test]
393 fn test_from() {
394 Python::with_gil(|py| {
395 let src = b"Hello Python";
396 let bytearray = PyByteArray::new(py, src);
397
398 let ba: PyObject = bytearray.into();
399 let bytearray = PyByteArray::from(ba.bind(py)).unwrap();
400
401 assert_eq!(src, unsafe { bytearray.as_bytes() });
402 });
403 }
404
405 #[test]
406 fn test_from_err() {
407 Python::with_gil(|py| {
408 if let Err(err) = PyByteArray::from(py.None().bind(py)) {
409 assert!(err.is_instance_of::<exceptions::PyTypeError>(py));
410 } else {
411 panic!("error");
412 }
413 });
414 }
415
416 #[test]
417 fn test_try_from() {
418 Python::with_gil(|py| {
419 let src = b"Hello Python";
420 let bytearray: &Bound<'_, PyAny> = &PyByteArray::new(py, src);
421 let bytearray: Bound<'_, PyByteArray> = TryInto::try_into(bytearray).unwrap();
422
423 assert_eq!(src, unsafe { bytearray.as_bytes() });
424 });
425 }
426
427 #[test]
428 fn test_resize() {
429 Python::with_gil(|py| {
430 let src = b"Hello Python";
431 let bytearray = PyByteArray::new(py, src);
432
433 bytearray.resize(20).unwrap();
434 assert_eq!(20, bytearray.len());
435 });
436 }
437
438 #[test]
439 fn test_byte_array_new_with() -> super::PyResult<()> {
440 Python::with_gil(|py| -> super::PyResult<()> {
441 let py_bytearray = PyByteArray::new_with(py, 10, |b: &mut [u8]| {
442 b.copy_from_slice(b"Hello Rust");
443 Ok(())
444 })?;
445 let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
446 assert_eq!(bytearray, b"Hello Rust");
447 Ok(())
448 })
449 }
450
451 #[test]
452 fn test_byte_array_new_with_zero_initialised() -> super::PyResult<()> {
453 Python::with_gil(|py| -> super::PyResult<()> {
454 let py_bytearray = PyByteArray::new_with(py, 10, |_b: &mut [u8]| Ok(()))?;
455 let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
456 assert_eq!(bytearray, &[0; 10]);
457 Ok(())
458 })
459 }
460
461 #[test]
462 fn test_byte_array_new_with_error() {
463 use crate::exceptions::PyValueError;
464 Python::with_gil(|py| {
465 let py_bytearray_result = PyByteArray::new_with(py, 10, |_b: &mut [u8]| {
466 Err(PyValueError::new_err("Hello Crustaceans!"))
467 });
468 assert!(py_bytearray_result.is_err());
469 assert!(py_bytearray_result
470 .err()
471 .unwrap()
472 .is_instance_of::<PyValueError>(py));
473 })
474 }
475}
476