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!(&section, 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}