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