Added Rust initialisation of Python-allocated bytes (#1074)

* Added Rust initialisation of Python-allocated bytes

* Added unsafe PyBytes::new_with_unit constructor

* Added examples to PyBytes::new_with and PyBytes::new_with_uninit (now with MaybeUninit<u8>)

* Fixed doc test imports for PyBytes::new_with and PyBytes::new_with_uninit

* Fixed clippy error in PyBytes::new_with_uninit test

* Added PyByteArray::new_with and Removed PyBytes::new_with_uninit

* Small doc fixes + FnOnce init closure for PyBytes::new_with and PyByteArray::new_with

Co-authored-by: David Hewitt <1939362+davidhewitt@users.noreply.github.com>

* Fixed where clause formatting in PyBytes::new_with and PyByteArray::new_with

Co-authored-by: David Hewitt <1939362+davidhewitt@users.noreply.github.com>
This commit is contained in:
Moritz Langenstein 2020-08-03 09:04:06 +01:00 committed by GitHub
parent 80224207cc
commit a877300f72
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 112 additions and 0 deletions

View File

@ -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)

View File

@ -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]);
}
}

View File

@ -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]);
}
}