diff --git a/src/types/list.rs b/src/types/list.rs index 4a0586e4fcd..406d4d82aca 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -178,26 +178,15 @@ where } } -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 IntoPy for [T; N] +where + T: ToPyObject, +{ + fn into_py(self, py: Python) -> PyObject { + self.as_ref().to_object(py) } } -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 -); - impl ToPyObject for Vec where T: ToPyObject, diff --git a/src/types/mod.rs b/src/types/mod.rs index 96459f37dc2..34c64e3f9f6 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -244,3 +244,37 @@ mod slice; mod string; mod tuple; mod typeobject; + +struct ArrayGuard { + dst: *mut T, + initialized: usize, +} + +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); + unsafe { + core::ptr::drop_in_place(initialized_part); + } + } +} + +fn try_create_array(mut cb: F) -> Result<[T; N], E> +where + F: FnMut(usize) -> Result, +{ + let mut array: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit(); + let mut guard: ArrayGuard = ArrayGuard { + dst: array.as_mut_ptr() as _, + 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; + } + core::mem::forget(guard); + Ok(array.assume_init()) + } +} diff --git a/src/types/sequence.rs b/src/types/sequence.rs index 423021d02a3..c6e8eb1261b 100644 --- a/src/types/sequence.rs +++ b/src/types/sequence.rs @@ -257,59 +257,39 @@ impl PySequence { } } -macro_rules! array_impls { - ($($N:expr),+) => { - $( - impl<'a, T> FromPyObject<'a> for [T; $N] - where - T: Copy + Default + FromPyObject<'a>, - { - #[cfg(not(feature = "nightly"))] - fn extract(obj: &'a PyAny) -> PyResult { - let mut array = [T::default(); $N]; - extract_sequence_into_slice(obj, &mut array)?; - Ok(array) - } +impl<'a, T, const N: usize> FromPyObject<'a> for [T; N] +where + T: FromPyObject<'a>, +{ + #[cfg(not(feature = "nightly"))] + fn extract(obj: &'a PyAny) -> PyResult { + create_array_from_obj(obj) + } - #[cfg(feature = "nightly")] - default fn extract(obj: &'a PyAny) -> PyResult { - let mut array = [T::default(); $N]; - extract_sequence_into_slice(obj, &mut array)?; - Ok(array) - } - } + #[cfg(feature = "nightly")] + default fn extract(obj: &'a PyAny) -> PyResult { + create_array_from_obj(obj) + } +} - #[cfg(feature = "nightly")] - impl<'source, T> FromPyObject<'source> for [T; $N] - where - for<'a> T: Default + FromPyObject<'a> + crate::buffer::Element, - { - fn extract(obj: &'source PyAny) -> PyResult { - let mut array = [T::default(); $N]; - // first try buffer protocol - if unsafe { 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()); - return Ok(array); - } - buf.release(obj.py()); - } - } - // fall back to sequence protocol - extract_sequence_into_slice(obj, &mut array)?; - Ok(array) - } +#[cfg(feature = "nightly")] +impl<'source, T, const N: usize> FromPyObject<'source> for [T; N] +where + for<'a> T: FromPyObject<'a> + crate::buffer::Element, +{ + fn extract(obj: &'source PyAny) -> PyResult { + let mut array = create_array_from_obj(obj)?; + 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()); + return Ok(array); } - )+ + buf.release(obj.py()); + } + Ok(array) } } -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 -); - impl<'a, T> FromPyObject<'a> for Vec where T: FromPyObject<'a>, @@ -345,32 +325,30 @@ where } } -fn extract_sequence<'s, T>(obj: &'s PyAny) -> PyResult> +fn create_array_from_obj<'s, T, const N: usize>(obj: &'s PyAny) -> PyResult<[T; N]> where T: FromPyObject<'s>, { let seq = ::try_from(obj)?; - let mut v = Vec::with_capacity(seq.len().unwrap_or(0) as usize); - for item in seq.iter()? { - v.push(item?.extract::()?); - } - Ok(v) + crate::types::try_create_array(|idx| { + seq.get_item(idx as isize) + .map_err(|_| { + exceptions::PyBufferError::new_err("Slice length does not match buffer length.") + })? + .extract::() + }) } -fn extract_sequence_into_slice<'s, T>(obj: &'s PyAny, slice: &mut [T]) -> PyResult<()> +fn extract_sequence<'s, T>(obj: &'s PyAny) -> PyResult> where T: FromPyObject<'s>, { let seq = ::try_from(obj)?; - if seq.len()? as usize != slice.len() { - return Err(exceptions::PyBufferError::new_err( - "Slice length does not match buffer length.", - )); - } - for (value, item) in slice.iter_mut().zip(seq.iter()?) { - *value = item?.extract::()?; + let mut v = Vec::with_capacity(seq.len().unwrap_or(0) as usize); + for item in seq.iter()? { + v.push(item?.extract::()?); } - Ok(()) + Ok(v) } impl<'v> PyTryFrom<'v> for PySequence {