From 7ead166d9dbc5a5f0a96014a0df1282a7002b83b Mon Sep 17 00:00:00 2001 From: David Hewitt <1939362+davidhewitt@users.noreply.github.com> Date: Sat, 24 Apr 2021 08:42:44 +0100 Subject: [PATCH] array: safer implementation of try_create_array --- CHANGELOG.md | 2 +- src/conversions/array.rs | 167 +++++++++++++++++++-------------------- src/lib.rs | 1 - src/utils.rs | 8 -- 4 files changed, 83 insertions(+), 95 deletions(-) delete mode 100644 src/utils.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c6fb353641..4bb32448f8b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,7 +13,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. - Support PyPy 3.7. [#1538](https://github.com/PyO3/pyo3/pull/1538) ### Added -- Add conversion for `[T; N]` for all `N` on Rust 1.51 and up. [#1128](https://github.com/PyO3/pyo3/pull/1128) +- Add conversions for `[T; N]` for all `N` on Rust 1.51 and up. [#1128](https://github.com/PyO3/pyo3/pull/1128) - Add conversions between `OsStr`/`OsString`/`Path`/`PathBuf` and Python strings. [#1379](https://github.com/PyO3/pyo3/pull/1379) - Add `#[pyo3(from_py_with = "...")]` attribute for function arguments and struct fields to override the default from-Python conversion. [#1411](https://github.com/PyO3/pyo3/pull/1411) - Add FFI definition `PyCFunction_CheckExact` for Python 3.9 and later. [#1425](https://github.com/PyO3/pyo3/pull/1425) diff --git a/src/conversions/array.rs b/src/conversions/array.rs index 8f6f321d9b9..3071d147cd3 100644 --- a/src/conversions/array.rs +++ b/src/conversions/array.rs @@ -1,9 +1,21 @@ -use crate::{FromPyObject, IntoPy, PyAny, PyObject, PyResult, PyTryFrom, Python, ToPyObject}; +use crate::{ + exceptions, FromPyObject, IntoPy, PyAny, PyErr, PyObject, PyResult, PyTryFrom, Python, + ToPyObject, +}; #[cfg(not(min_const_generics))] macro_rules! array_impls { ($($N:expr),+) => { $( + impl IntoPy for [T; $N] + where + T: ToPyObject + { + fn into_py(self, py: Python) -> PyObject { + self.as_ref().to_object(py) + } + } + impl<'a, T> FromPyObject<'a> for [T; $N] where T: Copy + Default + FromPyObject<'a>, @@ -55,6 +67,16 @@ array_impls!( 26, 27, 28, 29, 30, 31, 32 ); +#[cfg(min_const_generics)] +impl IntoPy for [T; N] +where + T: ToPyObject, +{ + fn into_py(self, py: Python) -> PyObject { + self.as_ref().to_object(py) + } +} + #[cfg(min_const_generics)] impl<'a, T, const N: usize> FromPyObject<'a> for [T; N] where @@ -71,60 +93,27 @@ where } } -#[cfg(not(min_const_generics))] -macro_rules! array_impls { - ($($N:expr),+) => { - $( - impl IntoPy for [T; $N] - where - T: ToPyObject - { - fn into_py(self, py: Python) -> PyObject { - self.as_ref().to_object(py) - } - } - )+ - } -} - -#[cfg(not(min_const_generics))] -array_impls!( - 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, - 26, 27, 28, 29, 30, 31, 32 -); - -#[cfg(min_const_generics)] -impl IntoPy for [T; N] -where - T: ToPyObject, -{ - fn into_py(self, py: Python) -> PyObject { - self.as_ref().to_object(py) - } -} - #[cfg(all(min_const_generics, feature = "nightly"))] impl<'source, T, const N: usize> FromPyObject<'source> for [T; N] where - for<'a> T: FromPyObject<'a> + crate::buffer::Element, + for<'a> T: Default + FromPyObject<'a> + crate::buffer::Element, { fn extract(obj: &'source PyAny) -> PyResult { - let mut array: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit(); + use crate::{AsPyPointer, PyNativeType}; + let mut array = [T::default(); N]; // first try buffer protocol if unsafe { crate::ffi::PyObject_CheckBuffer(obj.as_ptr()) } == 1 { if let Ok(buf) = crate::buffer::PyBuffer::get(obj) { if buf.dimensions() == 1 && buf.copy_to_slice(obj.py(), &mut array).is_ok() { buf.release(obj.py()); - // SAFETY: The array should be fully filled by `copy_to_slice` - return Ok(unsafe { array.assume_init() }); + return Ok(array); } buf.release(obj.py()); } } // fall back to sequence protocol _extract_sequence_into_slice(obj, &mut array)?; - // SAFETY: The array should be fully filled by `_extract_sequence_into_slice` - Ok(unsafe { array.assume_init() }) + Ok(array) } } @@ -135,102 +124,110 @@ where { let seq = ::try_from(obj)?; let expected_len = seq.len()? as usize; - let mut counter = 0; - try_create_array(&mut counter, |idx| { + array_try_from_fn(|idx| { seq.get_item(idx as isize) - .map_err(|_| crate::utils::invalid_sequence_length(expected_len, idx + 1))? + .map_err(|_| invalid_sequence_length(expected_len, idx + 1))? .extract::() }) } -fn _extract_sequence_into_slice<'s, T>(obj: &'s PyAny, slice: &mut [T]) -> PyResult<()> -where - T: FromPyObject<'s>, -{ - let seq = ::try_from(obj)?; - let expected_len = seq.len()? as usize; - if expected_len != slice.len() { - return Err(crate::utils::invalid_sequence_length( - expected_len, - slice.len(), - )); - } - for (value, item) in slice.iter_mut().zip(seq.iter()?) { - *value = item?.extract::()?; - } - Ok(()) -} - +// TODO use std::array::try_from_fn, if that stabilises: +// (https://github.com/rust-lang/rust/pull/75644) #[cfg(min_const_generics)] -fn try_create_array(counter: &mut usize, mut cb: F) -> Result<[T; N], E> +fn array_try_from_fn(mut cb: F) -> Result<[T; N], E> where F: FnMut(usize) -> Result, { // Helper to safely create arrays since the standard library doesn't // provide one yet. Shouldn't be necessary in the future. - struct ArrayGuard<'a, T, const N: usize> { + struct ArrayGuard { dst: *mut T, - initialized: &'a mut usize, + initialized: usize, } - impl Drop for ArrayGuard<'_, T, N> { + impl Drop for ArrayGuard { fn drop(&mut self) { - debug_assert!(*self.initialized <= N); - let initialized_part = core::ptr::slice_from_raw_parts_mut(self.dst, *self.initialized); + debug_assert!(self.initialized <= N); + let initialized_part = core::ptr::slice_from_raw_parts_mut(self.dst, self.initialized); unsafe { core::ptr::drop_in_place(initialized_part); } } } + // [MaybeUninit; N] would be "nicer" but is actually difficult to create - there are nightly + // APIs which would make this easier. let mut array: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit(); - let guard: ArrayGuard = ArrayGuard { + let mut guard: ArrayGuard = ArrayGuard { dst: array.as_mut_ptr() as _, - initialized: counter, + initialized: 0, }; unsafe { - for (idx, value_ptr) in (&mut *array.as_mut_ptr()).iter_mut().enumerate() { - core::ptr::write(value_ptr, cb(idx)?); - *guard.initialized += 1; + let mut value_ptr = array.as_mut_ptr() as *mut T; + for i in 0..N { + core::ptr::write(value_ptr, cb(i)?); + value_ptr = value_ptr.offset(1); + guard.initialized += 1; } core::mem::forget(guard); Ok(array.assume_init()) } } +fn _extract_sequence_into_slice<'s, T>(obj: &'s PyAny, slice: &mut [T]) -> PyResult<()> +where + T: FromPyObject<'s>, +{ + let seq = ::try_from(obj)?; + let expected_len = seq.len()? as usize; + if expected_len != slice.len() { + return Err(invalid_sequence_length(expected_len, slice.len())); + } + for (value, item) in slice.iter_mut().zip(seq.iter()?) { + *value = item?.extract::()?; + } + Ok(()) +} + +pub fn invalid_sequence_length(expected: usize, actual: usize) -> PyErr { + exceptions::PyValueError::new_err(format!( + "expected a sequence of length {} (got {})", + expected, actual + )) +} + #[cfg(test)] mod test { use crate::Python; #[cfg(min_const_generics)] use std::{ panic, - sync::{Arc, Mutex}, - thread::sleep, - time, + sync::atomic::{AtomicUsize, Ordering}, }; #[cfg(min_const_generics)] #[test] - fn try_create_array() { - #[allow(clippy::mutex_atomic)] - let counter = Arc::new(Mutex::new(0)); - let counter_unwind = Arc::clone(&counter); + fn array_try_from_fn() { + static DROP_COUNTER: AtomicUsize = AtomicUsize::new(0); + struct CountDrop; + impl Drop for CountDrop { + fn drop(&mut self) { + DROP_COUNTER.fetch_add(1, Ordering::SeqCst); + } + } let _ = catch_unwind_silent(move || { - let mut locked = counter_unwind.lock().unwrap(); - let _: Result<[i32; 4], _> = super::try_create_array(&mut *locked, |idx| { + let _: Result<[CountDrop; 4], ()> = super::array_try_from_fn(|idx| { if idx == 2 { panic!("peek a boo"); } - Ok::<_, ()>(1) + Ok(CountDrop) }); }); - sleep(time::Duration::from_secs(2)); - assert_eq!(*counter.lock().unwrap_err().into_inner(), 2); + assert_eq!(DROP_COUNTER.load(Ordering::SeqCst), 2); } - #[cfg(not(min_const_generics))] #[test] - fn test_extract_bytearray_to_array() { + fn test_extract_small_bytearray_to_array() { let gil = Python::acquire_gil(); let py = gil.python(); let v: [u8; 3] = py diff --git a/src/lib.rs b/src/lib.rs index f33759105a4..69735f74c24 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -186,7 +186,6 @@ pub mod pyclass_slots; mod python; pub mod type_object; pub mod types; -mod utils; #[cfg(feature = "serde")] pub mod serde; diff --git a/src/utils.rs b/src/utils.rs deleted file mode 100644 index 82c629fc9fc..00000000000 --- a/src/utils.rs +++ /dev/null @@ -1,8 +0,0 @@ -use crate::{exceptions, PyErr}; - -pub fn invalid_sequence_length(expected: usize, actual: usize) -> PyErr { - exceptions::PyValueError::new_err(format!( - "expected a sequence of length {} (got {})", - expected, actual - )) -}