Support arbitrary arrays

Author: c410-f3r

src/lib.rs

@@ -1,4 +1,5 @@
-#![cfg_attr(feature = "nightly", feature(specialization))]
+#![cfg_attr(feature = "nightly", allow(incomplete_features))]
+#![cfg_attr(feature = "nightly", feature(min_const_generics, specialization))]
 #![allow(clippy::missing_safety_doc)] // FIXME (#698)
 
 //! Rust bindings to the Python interpreter.

src/types/list.rs

@@ -178,6 +178,17 @@ where
     }
 }
 
+#[cfg(feature = "nightly")]
+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(not(feature = "nightly"))]
 macro_rules! array_impls {
     ($($N:expr),+) => {
         $(
@@ -193,6 +204,7 @@ macro_rules! array_impls {
     }
 }
 
+#[cfg(not(feature = "nightly"))]
 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

src/types/mod.rs

@@ -235,3 +235,40 @@ mod slice;
 mod string;
 mod tuple;
 mod typeobject;
+
+#[cfg(feature = "nightly")]
+struct ArrayGuard<T, const N: usize> {
+    dst: *mut T,
+    initialized: usize,
+}
+
+#[cfg(feature = "nightly")]
+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);
+        unsafe {
+            core::ptr::drop_in_place(initialized_part);
+        }
+    }
+}
+
+#[cfg(feature = "nightly")]
+fn create_array<F, T, const N: usize>(mut cb: F) -> [T; N]
+where
+    F: FnMut(usize) -> T,
+{
+    let mut array: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit();
+    let mut guard: ArrayGuard<T, N> = 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);
+        array.assume_init()
+    }
+}

src/types/sequence.rs

@@ -257,6 +257,40 @@ impl PySequence {
     }
 }
 
+#[cfg(feature = "nightly")]
+impl<'a, T, const N: usize> FromPyObject<'a> for [T; N]
+where
+    T: Copy + Default + FromPyObject<'a>,
+{
+    default fn extract(obj: &'a PyAny) -> PyResult<Self> {
+        let mut array = crate::types::create_array(|_| T::default());
+        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: Default + FromPyObject<'a> + crate::buffer::Element,
+{
+    fn extract(obj: &'source PyAny) -> PyResult<Self> {
+        let mut array = crate::types::create_array(|_| T::default());
+        // first try buffer protocol
+        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(not(feature = "nightly"))]
 macro_rules! array_impls {
     ($($N:expr),+) => {
         $(
@@ -303,6 +337,7 @@ macro_rules! array_impls {
     }
 }
 
+#[cfg(not(feature = "nightly"))]
 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