luxolus/pyo3
2
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

Merge pull request #1083 from MoritzLangenstein/fallible_py_bytes_bytearray_new_with

Browse source 
Changed PyByte::new_init and PyByteArray::new_init such that init can fail
This commit is contained in:
Yuji Kanagawa 2020-08-13 13:18:58 +09:00 committed by GitHub
parent 565e36d989 e6dc4b2d4d
commit 9ab7225ee0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 77 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
src/types/bytearray.rs
View file

@ -1,7 +1,7 @@
// Copyright (c) 2017-present PyO3 Project and Contributors
use crate::err::{PyErr, PyResult};
use crate::instance::PyNativeType;
use crate::{ffi, AsPyPointer, PyAny, Python};
use crate::{ffi, AsPyPointer, Py, PyAny, Python};
use std::os::raw::c_char;
use std::slice;
@ -23,36 +23,40 @@ impl PyByteArray {
/// 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.
/// * If Python raises a MemoryError on the allocation, `new_with` will return
/// it inside `Err`.
/// * If `init` returns `Err(e)`, `new_with` will return `Err(e)`.
/// * If `init` returns `Ok(())`, `new_with` will return `Ok(&PyByteArray)`.
///
/// # Example
/// ```
/// use pyo3::{prelude::*, types::PyByteArray};
/// Python::with_gil(|py| {
/// Python::with_gil(|py| -> PyResult<()> {
/// let py_bytearray = PyByteArray::new_with(py, 10, |bytes: &mut [u8]| {
/// bytes.copy_from_slice(b"Hello Rust");
/// });
/// Ok(())
/// })?;
/// let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
/// assert_eq!(bytearray, b"Hello Rust");
/// Ok(())
/// });
/// ```
pub fn new_with<F>(py: Python, len: usize, init: F) -> &PyByteArray
pub fn new_with<F>(py: Python, len: usize, init: F) -> PyResult<&PyByteArray>
where
F: FnOnce(&mut [u8]),
F: FnOnce(&mut [u8]) -> PyResult<()>,
{
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 pyptr =
ffi::PyByteArray_FromStringAndSize(std::ptr::null(), len as ffi::Py_ssize_t);
// Check for an allocation error and return it
let pypybytearray: Py<PyByteArray> = Py::from_owned_ptr_or_err(py, 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
// If init returns an Err, pypybytearray will automatically deallocate the buffer
init(std::slice::from_raw_parts_mut(buffer, len)).map(|_| pypybytearray.into_ref(py))
}
}
@ -263,22 +267,40 @@ mod test {
}
#[test]
fn test_byte_array_new_with() {
fn test_byte_array_new_with() -> super::PyResult<()> {
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");
});
Ok(())
})?;
let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
assert_eq!(bytearray, b"Hello Rust");
Ok(())
}
#[test]
fn test_byte_array_new_with_zero_initialised() {
fn test_byte_array_new_with_zero_initialised() -> super::PyResult<()> {
let gil = Python::acquire_gil();
let py = gil.python();
let py_bytearray = PyByteArray::new_with(py, 10, |_b: &mut [u8]| ());
let py_bytearray = PyByteArray::new_with(py, 10, |_b: &mut [u8]| Ok(()))?;
let bytearray: &[u8] = unsafe { py_bytearray.as_bytes() };
assert_eq!(bytearray, &[0; 10]);
Ok(())
}
#[test]
fn test_byte_array_new_with_error() {
use crate::exceptions::PyValueError;
let gil = Python::acquire_gil();
let py = gil.python();
let py_bytearray_result = PyByteArray::new_with(py, 10, |_b: &mut [u8]| {
Err(PyValueError::py_err("Hello Crustaceans!"))
});
assert!(py_bytearray_result.is_err());
assert!(py_bytearray_result
.err()
.unwrap()
.is_instance::<PyValueError>(py));
}
}

56
src/types/bytes.rs
View file

@ -1,5 +1,5 @@
use crate::{
ffi, AsPyPointer, FromPyObject, IntoPy, PyAny, PyObject, PyResult, PyTryFrom, Python,
ffi, AsPyPointer, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, PyTryFrom, Python,
ToPyObject,
};
use std::ops::Index;
@ -28,8 +28,10 @@ impl PyBytes {
/// 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.
/// * If Python raises a MemoryError on the allocation, `new_with` will return
/// it inside `Err`.
/// * If `init` returns `Err(e)`, `new_with` will return `Err(e)`.
/// * If `init` returns `Ok(())`, `new_with` will return `Ok(&PyBytes)`.
///
/// # Example
/// ```
@ -37,28 +39,28 @@ impl PyBytes {
/// Python::with_gil(|py| -> PyResult<()> {
/// let py_bytes = PyBytes::new_with(py, 10, |bytes: &mut [u8]| {
/// bytes.copy_from_slice(b"Hello Rust");
/// });
/// Ok(())
/// })?;
/// 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
pub fn new_with<F>(py: Python, len: usize, init: F) -> PyResult<&PyBytes>
where
F: FnOnce(&mut [u8]),
F: FnOnce(&mut [u8]) -> PyResult<()>,
{
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 pyptr = ffi::PyBytes_FromStringAndSize(std::ptr::null(), len as ffi::Py_ssize_t);
// Check for an allocation error and return it
let pypybytes: Py<PyBytes> = 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
init(std::slice::from_raw_parts_mut(buffer, len));
pybytes
// If init returns an Err, pypybytearray will automatically deallocate the buffer
init(std::slice::from_raw_parts_mut(buffer, len)).map(|_| pypybytes.into_ref(py))
}
}
@ -129,22 +131,40 @@ mod test {
}
#[test]
fn test_bytes_new_with() {
fn test_bytes_new_with() -> super::PyResult<()> {
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();
Ok(())
})?;
let bytes: &[u8] = FromPyObject::extract(py_bytes)?;
assert_eq!(bytes, b"Hello Rust");
Ok(())
}
#[test]
fn test_bytes_new_with_zero_initialised() {
fn test_bytes_new_with_zero_initialised() -> super::PyResult<()> {
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();
let py_bytes = PyBytes::new_with(py, 10, |_b: &mut [u8]| Ok(()))?;
let bytes: &[u8] = FromPyObject::extract(py_bytes)?;
assert_eq!(bytes, &[0; 10]);
Ok(())
}
#[test]
fn test_bytes_new_with_error() {
use crate::exceptions::PyValueError;
let gil = Python::acquire_gil();
let py = gil.python();
let py_bytes_result = PyBytes::new_with(py, 10, |_b: &mut [u8]| {
Err(PyValueError::py_err("Hello Crustaceans!"))
});
assert!(py_bytes_result.is_err());
assert!(py_bytes_result
.err()
.unwrap()
.is_instance::<PyValueError>(py));
}
}