diff --git a/CHANGELOG.md b/CHANGELOG.md index ab01c9cf69d..0b3bddbabe6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. - Add `Python::with_gil` for executing a closure with the Python GIL. [#1037](https://github.com/PyO3/pyo3/pull/1037) - Implement `Debug` for `PyIterator`. [#1051](https://github.com/PyO3/pyo3/pull/1051) - Implement type information for conversion failures. [#1050](https://github.com/PyO3/pyo3/pull/1050) +- Add `PyBytes::new_with` and `PyByteArray::new_with` for initialising Python-allocated bytes and bytearrays using a closure. [#1074](https://github.com/PyO3/pyo3/pull/1074) ### Changed - Exception types have been renamed from e.g. `RuntimeError` to `PyRuntimeError`, and are now only accessible by `&T` or `Py` similar to other Python-native types. The old names continue to exist but are deprecated. [#1024](https://github.com/PyO3/pyo3/pull/1024) diff --git a/src/types/bytearray.rs b/src/types/bytearray.rs index 5ddbbee5dd0..71a2ded98d7 100644 --- a/src/types/bytearray.rs +++ b/src/types/bytearray.rs @@ -21,6 +21,41 @@ impl PyByteArray { unsafe { py.from_owned_ptr::(ffi::PyByteArray_FromStringAndSize(ptr, len)) } } + /// Creates a new Python `bytearray` object with an `init` closure to write its contents. + /// Before calling `init` the bytearray is zero-initialised. + /// + /// Panics if out of memory. + /// + /// # Example + /// ``` + /// use pyo3::{prelude::*, types::PyByteArray}; + /// Python::with_gil(|py| { + /// let py_bytearray = PyByteArray::new_with(py, 10, |bytes: &mut [u8]| { + /// bytes.copy_from_slice(b"Hello Rust"); + /// }); + /// let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() }; + /// assert_eq!(bytearray, b"Hello Rust"); + /// }); + /// ``` + pub fn new_with(py: Python, len: usize, init: F) -> &PyByteArray + where + F: FnOnce(&mut [u8]), + { + unsafe { + let length = len as ffi::Py_ssize_t; + let pyptr = ffi::PyByteArray_FromStringAndSize(std::ptr::null(), length); + // Iff pyptr is null, py.from_owned_ptr(pyptr) will panic + let pybytearray = py.from_owned_ptr(pyptr); + let buffer = ffi::PyByteArray_AsString(pyptr) as *mut u8; + debug_assert!(!buffer.is_null()); + // Zero-initialise the uninitialised bytearray + std::ptr::write_bytes(buffer, 0u8, len); + // (Further) Initialise the bytearray in init + init(std::slice::from_raw_parts_mut(buffer, len)); + pybytearray + } + } + /// Creates a new Python bytearray object from another PyObject that /// implements the buffer protocol. pub fn from<'p, I>(py: Python<'p>, src: &'p I) -> PyResult<&'p PyByteArray> @@ -227,4 +262,24 @@ mod test { bytearray.resize(20).unwrap(); assert_eq!(20, bytearray.len()); } + + #[test] + fn test_byte_array_new_with() { + let gil = Python::acquire_gil(); + let py = gil.python(); + let py_bytearray = PyByteArray::new_with(py, 10, |b: &mut [u8]| { + b.copy_from_slice(b"Hello Rust"); + }); + let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() }; + assert_eq!(bytearray, b"Hello Rust"); + } + + #[test] + fn test_byte_array_new_with_zero_initialised() { + let gil = Python::acquire_gil(); + let py = gil.python(); + let py_bytearray = PyByteArray::new_with(py, 10, |_b: &mut [u8]| ()); + let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() }; + assert_eq!(bytearray, &[0; 10]); + } } diff --git a/src/types/bytes.rs b/src/types/bytes.rs index ba2dceab51a..8f8d3f76356 100644 --- a/src/types/bytes.rs +++ b/src/types/bytes.rs @@ -26,6 +26,42 @@ impl PyBytes { unsafe { py.from_owned_ptr(ffi::PyBytes_FromStringAndSize(ptr, len)) } } + /// Creates a new Python `bytes` object with an `init` closure to write its contents. + /// Before calling `init` the bytes' contents are zero-initialised. + /// + /// Panics if out of memory. + /// + /// # Example + /// ``` + /// use pyo3::{prelude::*, types::PyBytes}; + /// Python::with_gil(|py| -> PyResult<()> { + /// let py_bytes = PyBytes::new_with(py, 10, |bytes: &mut [u8]| { + /// bytes.copy_from_slice(b"Hello Rust"); + /// }); + /// let bytes: &[u8] = FromPyObject::extract(py_bytes)?; + /// assert_eq!(bytes, b"Hello Rust"); + /// Ok(()) + /// }); + /// ``` + pub fn new_with(py: Python, len: usize, init: F) -> &PyBytes + where + F: FnOnce(&mut [u8]), + { + unsafe { + let length = len as ffi::Py_ssize_t; + let pyptr = ffi::PyBytes_FromStringAndSize(std::ptr::null(), length); + // Iff pyptr is null, py.from_owned_ptr(pyptr) will panic + let pybytes = py.from_owned_ptr(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 + init(std::slice::from_raw_parts_mut(buffer, len)); + pybytes + } + } + /// Creates a new Python byte string object from a raw pointer and length. /// /// Panics if out of memory. @@ -91,4 +127,24 @@ mod test { let bytes = PyBytes::new(py, b"Hello World"); assert_eq!(bytes[1], b'e'); } + + #[test] + fn test_bytes_new_with() { + let gil = Python::acquire_gil(); + let py = gil.python(); + let py_bytes = PyBytes::new_with(py, 10, |b: &mut [u8]| { + b.copy_from_slice(b"Hello Rust"); + }); + let bytes: &[u8] = FromPyObject::extract(py_bytes).unwrap(); + assert_eq!(bytes, b"Hello Rust"); + } + + #[test] + fn test_bytes_new_with_zero_initialised() { + let gil = Python::acquire_gil(); + let py = gil.python(); + let py_bytes = PyBytes::new_with(py, 10, |_b: &mut [u8]| ()); + let bytes: &[u8] = FromPyObject::extract(py_bytes).unwrap(); + assert_eq!(bytes, &[0; 10]); + } }