array: safer implementation of try_create_array

Author: davidhewitt

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)

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<T> IntoPy<PyObject> 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<T, const N: usize> IntoPy<PyObject> 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<T> IntoPy<PyObject> 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<T, const N: usize> IntoPy<PyObject> 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<Self> {
-        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 = <crate::types::PySequence as PyTryFrom>::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::<T>()
     })
 }
 
-fn _extract_sequence_into_slice<'s, T>(obj: &'s PyAny, slice: &mut [T]) -> PyResult<()>
-where
-    T: FromPyObject<'s>,
-{
-    let seq = <crate::types::PySequence as PyTryFrom>::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::<T>()?;
-    }
-    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<E, F, T, const N: usize>(counter: &mut usize, mut cb: F) -> Result<[T; N], E>
+fn array_try_from_fn<E, F, T, const N: usize>(mut cb: F) -> Result<[T; N], E>
 where
     F: FnMut(usize) -> Result<T, E>,
 {
     // 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<T, const N: usize> {
         dst: *mut T,
-        initialized: &'a mut usize,
+        initialized: usize,
     }
 
-    impl<T, const N: usize> Drop for ArrayGuard<'_, T, N> {
+    impl<T, const N: usize> Drop for ArrayGuard<T, N> {
         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<T>; 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<T, N> = ArrayGuard {
+    let mut guard: ArrayGuard<T, N> = 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 = <crate::types::PySequence as PyTryFrom>::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::<T>()?;
+    }
+    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

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;