pyo3/types/bytearray.rs
1use crate::err::{PyErr, PyResult};
2use crate::ffi_ptr_ext::FfiPtrExt;
3use crate::instance::{Borrowed, Bound};
4use crate::py_result_ext::PyResultExt;
5use crate::sync::with_critical_section;
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 .cast_into_unchecked()
32 }
33 }
34
35 /// Creates a new Python `bytearray` object with an `init` closure to write its contents.
36 /// Before calling `init` the bytearray is zero-initialised.
37 /// * If Python raises a MemoryError on the allocation, `new_with` will return
38 /// it inside `Err`.
39 /// * If `init` returns `Err(e)`, `new_with` will return `Err(e)`.
40 /// * If `init` returns `Ok(())`, `new_with` will return `Ok(&PyByteArray)`.
41 ///
42 /// # Examples
43 ///
44 /// ```
45 /// use pyo3::{prelude::*, types::PyByteArray};
46 ///
47 /// # fn main() -> PyResult<()> {
48 /// Python::attach(|py| -> PyResult<()> {
49 /// let py_bytearray = PyByteArray::new_with(py, 10, |bytes: &mut [u8]| {
50 /// bytes.copy_from_slice(b"Hello Rust");
51 /// Ok(())
52 /// })?;
53 /// let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
54 /// assert_eq!(bytearray, b"Hello Rust");
55 /// Ok(())
56 /// })
57 /// # }
58 /// ```
59 pub fn new_with<F>(py: Python<'_>, len: usize, init: F) -> PyResult<Bound<'_, PyByteArray>>
60 where
61 F: FnOnce(&mut [u8]) -> PyResult<()>,
62 {
63 unsafe {
64 // Allocate buffer and check for an error
65 let pybytearray: Bound<'_, Self> =
66 ffi::PyByteArray_FromStringAndSize(std::ptr::null(), len as ffi::Py_ssize_t)
67 .assume_owned_or_err(py)?
68 .cast_into_unchecked();
69
70 let buffer: *mut u8 = ffi::PyByteArray_AsString(pybytearray.as_ptr()).cast();
71 debug_assert!(!buffer.is_null());
72 // Zero-initialise the uninitialised bytearray
73 std::ptr::write_bytes(buffer, 0u8, len);
74 // (Further) Initialise the bytearray in init
75 // If init returns an Err, pypybytearray will automatically deallocate the buffer
76 init(std::slice::from_raw_parts_mut(buffer, len)).map(|_| pybytearray)
77 }
78 }
79
80 /// Creates a new Python `bytearray` object from another Python object that
81 /// implements the buffer protocol.
82 pub fn from<'py>(src: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyByteArray>> {
83 unsafe {
84 ffi::PyByteArray_FromObject(src.as_ptr())
85 .assume_owned_or_err(src.py())
86 .cast_into_unchecked()
87 }
88 }
89}
90
91/// Implementation of functionality for [`PyByteArray`].
92///
93/// These methods are defined for the `Bound<'py, PyByteArray>` smart pointer, so to use method call
94/// syntax these methods are separated into a trait, because stable Rust does not yet support
95/// `arbitrary_self_types`.
96#[doc(alias = "PyByteArray")]
97pub trait PyByteArrayMethods<'py>: crate::sealed::Sealed {
98 /// Gets the length of the bytearray.
99 fn len(&self) -> usize;
100
101 /// Checks if the bytearray is empty.
102 fn is_empty(&self) -> bool;
103
104 /// Gets the start of the buffer containing the contents of the bytearray.
105 ///
106 /// # Safety
107 ///
108 /// See the safety requirements of [`PyByteArrayMethods::as_bytes`] and [`PyByteArrayMethods::as_bytes_mut`].
109 fn data(&self) -> *mut u8;
110
111 /// Extracts a slice of the `ByteArray`'s entire buffer.
112 ///
113 /// # Safety
114 ///
115 /// Mutation of the `bytearray` invalidates the slice. If it is used afterwards, the behavior is
116 /// undefined.
117 ///
118 /// These mutations may occur in Python code as well as from Rust:
119 /// - Calling methods like [`PyByteArrayMethods::as_bytes_mut`] and [`PyByteArrayMethods::resize`] will
120 /// invalidate the slice.
121 /// - Actions like dropping objects or raising exceptions can invoke `__del__`methods or signal
122 /// handlers, which may execute arbitrary Python code. This means that if Python code has a
123 /// reference to the `bytearray` you cannot safely use the vast majority of PyO3's API whilst
124 /// using the slice.
125 ///
126 /// As a result, this slice should only be used for short-lived operations without executing any
127 /// Python code, such as copying into a Vec.
128 /// For free-threaded Python support see also [`with_critical_section`].
129 ///
130 /// # Examples
131 ///
132 /// ```rust
133 /// use pyo3::prelude::*;
134 /// use pyo3::exceptions::PyRuntimeError;
135 /// use pyo3::sync::with_critical_section;
136 /// use pyo3::types::PyByteArray;
137 ///
138 /// #[pyfunction]
139 /// fn a_valid_function(bytes: &Bound<'_, PyByteArray>) -> PyResult<()> {
140 /// let section = with_critical_section(bytes, || {
141 /// // SAFETY: We promise to not let the interpreter regain control over the bytearray
142 /// // or invoke any PyO3 APIs while using the slice.
143 /// let slice = unsafe { bytes.as_bytes() };
144 ///
145 /// // Copy only a section of `bytes` while avoiding
146 /// // `to_vec` which copies the entire thing.
147 /// slice.get(6..11)
148 /// .map(Vec::from)
149 /// .ok_or_else(|| PyRuntimeError::new_err("input is not long enough"))
150 /// })?;
151 ///
152 /// // Now we can do things with `section` and call PyO3 APIs again.
153 /// // ...
154 /// # assert_eq!(§ion, b"world");
155 ///
156 /// Ok(())
157 /// }
158 /// # fn main() -> PyResult<()> {
159 /// # Python::attach(|py| -> PyResult<()> {
160 /// # let fun = wrap_pyfunction!(a_valid_function, py)?;
161 /// # let locals = pyo3::types::PyDict::new(py);
162 /// # locals.set_item("a_valid_function", fun)?;
163 /// #
164 /// # py.run(pyo3::ffi::c_str!(
165 /// # r#"b = bytearray(b"hello world")
166 /// # a_valid_function(b)
167 /// #
168 /// # try:
169 /// # a_valid_function(bytearray())
170 /// # except RuntimeError as e:
171 /// # assert str(e) == 'input is not long enough'"#),
172 /// # None,
173 /// # Some(&locals),
174 /// # )?;
175 /// #
176 /// # Ok(())
177 /// # })
178 /// # }
179 /// ```
180 ///
181 /// # Incorrect usage
182 ///
183 /// The following `bug` function is unsound ⚠️
184 ///
185 /// ```rust,no_run
186 /// # use pyo3::prelude::*;
187 /// # use pyo3::types::PyByteArray;
188 ///
189 /// # #[allow(dead_code)]
190 /// #[pyfunction]
191 /// fn bug(py: Python<'_>, bytes: &Bound<'_, PyByteArray>) {
192 /// // No critical section is being used.
193 /// // This means that for no-gil Python another thread could be modifying the
194 /// // bytearray concurrently and thus invalidate `slice` any time.
195 /// let slice = unsafe { bytes.as_bytes() };
196 ///
197 /// // This explicitly yields control back to the Python interpreter...
198 /// // ...but it's not always this obvious. Many things do this implicitly.
199 /// py.detach(|| {
200 /// // Python code could be mutating through its handle to `bytes`,
201 /// // which makes reading it a data race, which is undefined behavior.
202 /// println!("{:?}", slice[0]);
203 /// });
204 ///
205 /// // Python code might have mutated it, so we can not rely on the slice
206 /// // remaining valid. As such this is also undefined behavior.
207 /// println!("{:?}", slice[0]);
208 /// }
209 /// ```
210 unsafe fn as_bytes(&self) -> &[u8];
211
212 /// Extracts a mutable slice of the `ByteArray`'s entire buffer.
213 ///
214 /// # Safety
215 ///
216 /// Any other accesses of the `bytearray`'s buffer invalidate the slice. If it is used
217 /// afterwards, the behavior is undefined. The safety requirements of [`PyByteArrayMethods::as_bytes`]
218 /// apply to this function as well.
219 #[allow(clippy::mut_from_ref)]
220 unsafe fn as_bytes_mut(&self) -> &mut [u8];
221
222 /// Copies the contents of the bytearray to a Rust vector.
223 ///
224 /// # Examples
225 ///
226 /// ```
227 /// # use pyo3::prelude::*;
228 /// # use pyo3::types::PyByteArray;
229 /// # Python::attach(|py| {
230 /// let bytearray = PyByteArray::new(py, b"Hello World.");
231 /// let mut copied_message = bytearray.to_vec();
232 /// assert_eq!(b"Hello World.", copied_message.as_slice());
233 ///
234 /// copied_message[11] = b'!';
235 /// assert_eq!(b"Hello World!", copied_message.as_slice());
236 ///
237 /// pyo3::py_run!(py, bytearray, "assert bytearray == b'Hello World.'");
238 /// # });
239 /// ```
240 fn to_vec(&self) -> Vec<u8>;
241
242 /// Resizes the bytearray object to the new length `len`.
243 ///
244 /// Note that this will invalidate any pointers obtained by [PyByteArrayMethods::data], as well as
245 /// any (unsafe) slices obtained from [PyByteArrayMethods::as_bytes] and [PyByteArrayMethods::as_bytes_mut].
246 fn resize(&self, len: usize) -> PyResult<()>;
247}
248
249impl<'py> PyByteArrayMethods<'py> for Bound<'py, PyByteArray> {
250 #[inline]
251 fn len(&self) -> usize {
252 // non-negative Py_ssize_t should always fit into Rust usize
253 unsafe { ffi::PyByteArray_Size(self.as_ptr()) as usize }
254 }
255
256 fn is_empty(&self) -> bool {
257 self.len() == 0
258 }
259
260 fn data(&self) -> *mut u8 {
261 self.as_borrowed().data()
262 }
263
264 unsafe fn as_bytes(&self) -> &[u8] {
265 unsafe { self.as_borrowed().as_bytes() }
266 }
267
268 #[allow(clippy::mut_from_ref)]
269 unsafe fn as_bytes_mut(&self) -> &mut [u8] {
270 unsafe { self.as_borrowed().as_bytes_mut() }
271 }
272
273 fn to_vec(&self) -> Vec<u8> {
274 with_critical_section(self, || {
275 // SAFETY:
276 // * `self` is a `Bound` object, which guarantees that the Python GIL is held.
277 // * For no-gil Python, a critical section is used in lieu of the GIL.
278 // * We don't interact with the interpreter
279 // * We don't mutate the underlying slice
280 unsafe { self.as_bytes() }.to_vec()
281 })
282 }
283
284 fn resize(&self, len: usize) -> PyResult<()> {
285 unsafe {
286 let result = ffi::PyByteArray_Resize(self.as_ptr(), len as ffi::Py_ssize_t);
287 if result == 0 {
288 Ok(())
289 } else {
290 Err(PyErr::fetch(self.py()))
291 }
292 }
293 }
294}
295
296impl<'a> Borrowed<'a, '_, PyByteArray> {
297 fn data(&self) -> *mut u8 {
298 unsafe { ffi::PyByteArray_AsString(self.as_ptr()).cast() }
299 }
300
301 #[allow(clippy::wrong_self_convention)]
302 unsafe fn as_bytes(self) -> &'a [u8] {
303 unsafe { slice::from_raw_parts(self.data(), self.len()) }
304 }
305
306 #[allow(clippy::wrong_self_convention)]
307 unsafe fn as_bytes_mut(self) -> &'a mut [u8] {
308 unsafe { slice::from_raw_parts_mut(self.data(), self.len()) }
309 }
310}
311
312impl<'py> TryFrom<&Bound<'py, PyAny>> for Bound<'py, PyByteArray> {
313 type Error = crate::PyErr;
314
315 /// Creates a new Python `bytearray` object from another Python object that
316 /// implements the buffer protocol.
317 fn try_from(value: &Bound<'py, PyAny>) -> Result<Self, Self::Error> {
318 PyByteArray::from(value)
319 }
320}
321
322#[cfg(test)]
323mod tests {
324 use crate::types::{PyAnyMethods, PyByteArray, PyByteArrayMethods};
325 use crate::{exceptions, Bound, Py, PyAny, Python};
326
327 #[test]
328 fn test_len() {
329 Python::attach(|py| {
330 let src = b"Hello Python";
331 let bytearray = PyByteArray::new(py, src);
332 assert_eq!(src.len(), bytearray.len());
333 });
334 }
335
336 #[test]
337 fn test_as_bytes() {
338 Python::attach(|py| {
339 let src = b"Hello Python";
340 let bytearray = PyByteArray::new(py, src);
341
342 let slice = unsafe { bytearray.as_bytes() };
343 assert_eq!(src, slice);
344 assert_eq!(bytearray.data() as *const _, slice.as_ptr());
345 });
346 }
347
348 #[test]
349 fn test_as_bytes_mut() {
350 Python::attach(|py| {
351 let src = b"Hello Python";
352 let bytearray = PyByteArray::new(py, src);
353
354 let slice = unsafe { bytearray.as_bytes_mut() };
355 assert_eq!(src, slice);
356 assert_eq!(bytearray.data(), slice.as_mut_ptr());
357
358 slice[0..5].copy_from_slice(b"Hi...");
359
360 assert_eq!(bytearray.str().unwrap(), "bytearray(b'Hi... Python')");
361 });
362 }
363
364 #[test]
365 fn test_to_vec() {
366 Python::attach(|py| {
367 let src = b"Hello Python";
368 let bytearray = PyByteArray::new(py, src);
369
370 let vec = bytearray.to_vec();
371 assert_eq!(src, vec.as_slice());
372 });
373 }
374
375 #[test]
376 fn test_from() {
377 Python::attach(|py| {
378 let src = b"Hello Python";
379 let bytearray = PyByteArray::new(py, src);
380
381 let ba: Py<PyAny> = bytearray.into();
382 let bytearray = PyByteArray::from(ba.bind(py)).unwrap();
383
384 assert_eq!(src, unsafe { bytearray.as_bytes() });
385 });
386 }
387
388 #[test]
389 fn test_from_err() {
390 Python::attach(|py| {
391 if let Err(err) = PyByteArray::from(py.None().bind(py)) {
392 assert!(err.is_instance_of::<exceptions::PyTypeError>(py));
393 } else {
394 panic!("error");
395 }
396 });
397 }
398
399 #[test]
400 fn test_try_from() {
401 Python::attach(|py| {
402 let src = b"Hello Python";
403 let bytearray: &Bound<'_, PyAny> = &PyByteArray::new(py, src);
404 let bytearray: Bound<'_, PyByteArray> = TryInto::try_into(bytearray).unwrap();
405
406 assert_eq!(src, unsafe { bytearray.as_bytes() });
407 });
408 }
409
410 #[test]
411 fn test_resize() {
412 Python::attach(|py| {
413 let src = b"Hello Python";
414 let bytearray = PyByteArray::new(py, src);
415
416 bytearray.resize(20).unwrap();
417 assert_eq!(20, bytearray.len());
418 });
419 }
420
421 #[test]
422 fn test_byte_array_new_with() -> super::PyResult<()> {
423 Python::attach(|py| -> super::PyResult<()> {
424 let py_bytearray = PyByteArray::new_with(py, 10, |b: &mut [u8]| {
425 b.copy_from_slice(b"Hello Rust");
426 Ok(())
427 })?;
428 let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
429 assert_eq!(bytearray, b"Hello Rust");
430 Ok(())
431 })
432 }
433
434 #[test]
435 fn test_byte_array_new_with_zero_initialised() -> super::PyResult<()> {
436 Python::attach(|py| -> super::PyResult<()> {
437 let py_bytearray = PyByteArray::new_with(py, 10, |_b: &mut [u8]| Ok(()))?;
438 let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
439 assert_eq!(bytearray, &[0; 10]);
440 Ok(())
441 })
442 }
443
444 #[test]
445 fn test_byte_array_new_with_error() {
446 use crate::exceptions::PyValueError;
447 Python::attach(|py| {
448 let py_bytearray_result = PyByteArray::new_with(py, 10, |_b: &mut [u8]| {
449 Err(PyValueError::new_err("Hello Crustaceans!"))
450 });
451 assert!(py_bytearray_result.is_err());
452 assert!(py_bytearray_result
453 .err()
454 .unwrap()
455 .is_instance_of::<PyValueError>(py));
456 })
457 }
458
459 // * wasm has no threading support
460 // * CPython 3.13t is unsound => test fails
461 #[cfg(all(
462 not(target_family = "wasm"),
463 any(Py_3_14, not(all(Py_3_13, Py_GIL_DISABLED)))
464 ))]
465 #[test]
466 fn test_data_integrity_in_critical_section() {
467 use crate::instance::Py;
468 use crate::sync::{with_critical_section, MutexExt};
469
470 use std::sync::atomic::{AtomicBool, Ordering};
471 use std::sync::Mutex;
472 use std::thread;
473 use std::thread::ScopedJoinHandle;
474 use std::time::Duration;
475
476 const SIZE: usize = 1_000_000;
477 const DATA_VALUE: u8 = 42;
478
479 fn make_byte_array(py: Python<'_>, size: usize, value: u8) -> Bound<'_, PyByteArray> {
480 PyByteArray::new_with(py, size, |b| {
481 b.fill(value);
482 Ok(())
483 })
484 .unwrap()
485 }
486
487 let data: Mutex<Py<PyByteArray>> = Mutex::new(Python::attach(|py| {
488 make_byte_array(py, SIZE, DATA_VALUE).unbind()
489 }));
490
491 fn get_data<'py>(
492 data: &Mutex<Py<PyByteArray>>,
493 py: Python<'py>,
494 ) -> Bound<'py, PyByteArray> {
495 data.lock_py_attached(py).unwrap().bind(py).clone()
496 }
497
498 fn set_data(data: &Mutex<Py<PyByteArray>>, new: Bound<'_, PyByteArray>) {
499 let py = new.py();
500 *data.lock_py_attached(py).unwrap() = new.unbind()
501 }
502
503 let running = AtomicBool::new(true);
504 let extending = AtomicBool::new(false);
505
506 // continuously extends and resets the bytearray in data
507 let worker1 = || {
508 let mut rounds = 0;
509 while running.load(Ordering::SeqCst) && rounds < 50 {
510 Python::attach(|py| {
511 let byte_array = get_data(&data, py);
512 extending.store(true, Ordering::SeqCst);
513 byte_array
514 .call_method("extend", (&byte_array,), None)
515 .unwrap();
516 extending.store(false, Ordering::SeqCst);
517 set_data(&data, make_byte_array(py, SIZE, DATA_VALUE));
518 rounds += 1;
519 });
520 }
521 };
522
523 // continuously checks the integrity of bytearray in data
524 let worker2 = || {
525 while running.load(Ordering::SeqCst) {
526 if !extending.load(Ordering::SeqCst) {
527 // wait until we have a chance to read inconsistent state
528 continue;
529 }
530 Python::attach(|py| {
531 let read = get_data(&data, py);
532 if read.len() == SIZE {
533 // extend is still not done => wait even more
534 return;
535 }
536 with_critical_section(&read, || {
537 // SAFETY: we are in a critical section
538 // This is the whole point of the test: make sure that a
539 // critical section is sufficient to ensure that the data
540 // read is consistent.
541 unsafe {
542 let bytes = read.as_bytes();
543 assert!(bytes.iter().rev().take(50).all(|v| *v == DATA_VALUE
544 && bytes.iter().take(50).all(|v| *v == DATA_VALUE)));
545 }
546 });
547 });
548 }
549 };
550
551 thread::scope(|s| {
552 let mut handle1 = Some(s.spawn(worker1));
553 let mut handle2 = Some(s.spawn(worker2));
554 let mut handles = [&mut handle1, &mut handle2];
555
556 let t0 = std::time::Instant::now();
557 while t0.elapsed() < Duration::from_secs(1) {
558 for handle in &mut handles {
559 if handle
560 .as_ref()
561 .map(ScopedJoinHandle::is_finished)
562 .unwrap_or(false)
563 {
564 let res = handle.take().unwrap().join();
565 if res.is_err() {
566 running.store(false, Ordering::SeqCst);
567 }
568 res.unwrap();
569 }
570 }
571 if handles.iter().any(|handle| handle.is_none()) {
572 break;
573 }
574 }
575 running.store(false, Ordering::SeqCst);
576 for handle in &mut handles {
577 if let Some(handle) = handle.take() {
578 handle.join().unwrap()
579 }
580 }
581 });
582 }
583}