Added Rust initialisation of Python-allocated bytes

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<T>` similar to other Python-native types. The old names continue to exist but are deprecated. [#1024](https://github.com/PyO3/pyo3/pull/1024)

src/types/bytearray.rs

@@ -21,6 +21,41 @@ impl PyByteArray {
         unsafe { py.from_owned_ptr::<PyByteArray>(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<F>(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]);
+    }
 }

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<F>(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]);
+    }
 }