diff --git a/src/types/bytes.rs b/src/types/bytes.rs index 448fe3b22c8..cdafcd3460c 100644 --- a/src/types/bytes.rs +++ b/src/types/bytes.rs @@ -1,7 +1,4 @@ -use crate::{ - ffi, AsPyPointer, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, PyTryFrom, Python, - ToPyObject, -}; +use crate::{ffi, AsPyPointer, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, PyTryFrom, Python, ToPyObject, PyErr}; use std::ops::Index; use std::os::raw::c_char; use std::slice::SliceIndex; @@ -64,6 +61,31 @@ impl PyBytes { } } + pub fn new_with_resize(py: Python, len: usize, init: F) -> PyResult<&PyBytes> + where + F: FnOnce(&mut [u8]) -> PyResult, + { + unsafe { + let pyptr = ffi::PyBytes_FromStringAndSize(std::ptr::null(), len as ffi::Py_ssize_t); + // Check for an allocation error and return it + let pypybytes: Py = Py::from_owned_ptr_or_err(py, pyptr)?; + let buffer = ffi::PyBytes_AsString(pyptr) as *mut u8; + debug_assert!(!buffer.is_null()); + // Zero-initialise the uninitialised bytestring + std::ptr::write_bytes(buffer, 0u8, len); + // (Further) Initialise the bytestring in init + // If init returns an Err, pypybytearray will automatically deallocate the buffer + let resize = init(std::slice::from_raw_parts_mut(buffer, len)).unwrap(); + let result = ffi::_PyBytes_Resize(&mut pypybytes.as_ptr(), resize as ffi::Py_ssize_t); + if result == 0 { + Ok(pypybytes.into_ref(py)) + } else { + Err(PyErr::fetch(py)) + } + + } + } + /// Creates a new Python byte string object from a raw pointer and length. /// /// Panics if out of memory.