WIP: Check dynamic type of individual elements to enable arrays holding Py<T> with custom types T.

Author: adamreichold

src/array.rs

@@ -127,8 +127,9 @@ impl<T, D> IntoPy<PyObject> for PyArray<T, D> {
 impl<'a, T: Element, D: Dimension> FromPyObject<'a> for &'a PyArray<T, D> {
     // here we do type-check three times
     // 1. Checks if the object is PyArray
-    // 2. Checks if the data type of the array is T
-    // 3. Checks if the dimension is same as D
+    // 2. Checks if the dimension is same as D
+    // 3. Checks if the data type of the array is T
+    // 4. Optionally checks if the elements of the array match T
     fn extract(ob: &'a PyAny) -> PyResult<Self> {
         let array = unsafe {
             if npyffi::PyArray_Check(ob.as_ptr()) == 0 {
@@ -137,19 +138,21 @@ impl<'a, T: Element, D: Dimension> FromPyObject<'a> for &'a PyArray<T, D> {
             &*(ob as *const PyAny as *const PyArray<T, D>)
         };
 
-        let src_dtype = array.dtype();
-        let dst_dtype = T::get_dtype(ob.py());
-        if !src_dtype.is_equiv_to(dst_dtype) {
-            return Err(TypeError::new(src_dtype, dst_dtype).into());
-        }
-
         let src_ndim = array.shape().len();
         if let Some(dst_ndim) = D::NDIM {
             if src_ndim != dst_ndim {
                 return Err(DimensionalityError::new(src_ndim, dst_ndim).into());
             }
         }
 
+        let src_dtype = array.dtype();
+        let dst_dtype = T::get_dtype(ob.py());
+        if !src_dtype.is_equiv_to(dst_dtype) {
+            return Err(TypeError::new(src_dtype, dst_dtype).into());
+        }
+
+        T::check_element_types(ob.py(), array)?;
+
         Ok(array)
     }
 }

src/dtype.rs

@@ -1,10 +1,17 @@
+use std::any::{type_name, TypeId};
 use std::mem::size_of;
 use std::os::raw::{c_int, c_long, c_longlong, c_short, c_uint, c_ulong, c_ulonglong, c_ushort};
 
+use ndarray::Dimension;
 use num_traits::{Bounded, Zero};
-use pyo3::{ffi, prelude::*, pyobject_native_type_core, types::PyType, AsPyPointer, PyNativeType};
+use pyo3::{
+    ffi, prelude::*, pyobject_native_type_core, type_object::PyTypeObject, types::PyType,
+    AsPyPointer, PyDowncastError, PyNativeType, PyTypeInfo,
+};
 
+use crate::array::PyArray;
 use crate::npyffi::{NpyTypes, PyArray_Descr, NPY_TYPES, PY_ARRAY_API};
+use crate::NpySingleIterBuilder;
 
 pub use num_complex::{Complex32, Complex64};
 
@@ -132,26 +139,12 @@ impl PyArrayDescr {
 /// #[pyclass]
 /// pub struct CustomElement;
 ///
-/// // The transparent wrapper is necessary as one cannot implement
-/// // a foreign trait (`Element`) on a foreign type (`Py`) directly.
-/// #[derive(Clone)]
-/// #[repr(transparent)]
-/// pub struct Wrapper(pub Py<CustomElement>);
-///
-/// unsafe impl Element for Wrapper {
-///     const IS_COPY: bool = false;
-///
-///     fn get_dtype(py: Python) -> &PyArrayDescr {
-///         PyArrayDescr::object(py)
-///     }
-/// }
-///
 /// Python::with_gil(|py| {
-///     let array = Array2::<Wrapper>::from_shape_fn((2, 3), |(_i, _j)| {
-///         Wrapper(Py::new(py, CustomElement).unwrap())
+///     let array = Array2::<Py<CustomElement>>::from_shape_fn((2, 3), |(_i, _j)| {
+///         Py::new(py, CustomElement).unwrap()
 ///     });
 ///
-///     let _array: &PyArray<Wrapper, _> = array.to_pyarray(py);
+///     let _array: &PyArray<Py<CustomElement>, _> = array.to_pyarray(py);
 /// });
 /// ```
 pub unsafe trait Element: Clone + Send {
@@ -164,6 +157,9 @@ pub unsafe trait Element: Clone + Send {
     /// that contain object-type fields.
     const IS_COPY: bool;
 
+    /// TODO
+    fn check_element_types<D: Dimension>(py: Python, array: &PyArray<Self, D>) -> PyResult<()>;
+
     /// Returns the associated array descriptor ("dtype") for the given type.
     fn get_dtype(py: Python) -> &PyArrayDescr;
 }
@@ -218,6 +214,12 @@ macro_rules! impl_element_scalar {
         $(#[$meta])*
         unsafe impl Element for $ty {
             const IS_COPY: bool = true;
+
+            fn check_element_types<D: Dimension>(_py: Python, _array: &PyArray<Self, D>) -> PyResult<()> {
+                // For scalar types, checking the dtype is sufficient.
+                Ok(())
+            }
+
             fn get_dtype(py: Python) -> &PyArrayDescr {
                 PyArrayDescr::from_npy_type(py, $npy_type)
             }
@@ -244,9 +246,30 @@ impl_element_scalar!(Complex64 => NPY_CDOUBLE,
 #[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
 impl_element_scalar!(usize, isize);
 
-unsafe impl Element for PyObject {
+unsafe impl<T> Element for Py<T>
+where
+    T: PyTypeInfo + 'static,
+{
     const IS_COPY: bool = false;
 
+    fn check_element_types<D: Dimension>(py: Python, array: &PyArray<Self, D>) -> PyResult<()> {
+        // `PyAny` can represent any Python object.
+        if TypeId::of::<PyAny>() == TypeId::of::<T>() {
+            return Ok(());
+        }
+
+        let type_object = T::type_object(py);
+        let iterator = NpySingleIterBuilder::readwrite(array).build()?;
+
+        for element in iterator {
+            if !type_object.is_instance(element)? {
+                return Err(PyDowncastError::new(array, type_name::<T>()).into());
+            }
+        }
+
+        Ok(())
+    }
+
     fn get_dtype(py: Python) -> &PyArrayDescr {
         PyArrayDescr::object(py)
     }