PyArrayDescr methods like in np.dtype + some missing stuff in npyffi

src/array.rs

@@ -1219,7 +1219,7 @@ impl<T: Element + AsPrimitive<f64>> PyArray<T, Ix1> {
                 start.as_(),
                 stop.as_(),
                 step.as_(),
-                T::get_dtype(py).get_typenum(),
+                T::get_dtype(py).num(),
             );
             Self::from_owned_ptr(py, ptr)
         }

src/dtype.rs

@@ -1,10 +1,21 @@
 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 std::os::raw::{
+    c_char, c_int, c_long, c_longlong, c_short, c_uint, c_ulong, c_ulonglong, c_ushort,
+};
 
 use num_traits::{Bounded, Zero};
-use pyo3::{ffi, prelude::*, pyobject_native_type_core, types::PyType, AsPyPointer, PyNativeType};
+use pyo3::{
+    ffi::{self, PyTuple_Size},
+    prelude::*,
+    pyobject_native_type_core,
+    types::{PyDict, PyTuple, PyType},
+    AsPyPointer, FromPyObject, FromPyPointer, PyNativeType,
+};
 
-use crate::npyffi::{NpyTypes, PyArray_Descr, NPY_TYPES, PY_ARRAY_API};
+use crate::npyffi::{
+    NpyTypes, PyArray_Descr, NPY_ALIGNED_STRUCT, NPY_BYTEORDER_CHAR, NPY_ITEM_HASOBJECT, NPY_TYPES,
+    PY_ARRAY_API,
+};
 
 pub use num_complex::{Complex32, Complex64};
 
@@ -57,21 +68,6 @@ impl PyArrayDescr {
         self.into_ptr() as _
     }
 
-    /// Returns the internal `PyType` that this `dtype` holds.
-    ///
-    /// # Example
-    /// ```
-    /// pyo3::Python::with_gil(|py| {
-    ///    let array = numpy::PyArray::from_vec(py, vec![0.0, 1.0, 2.0f64]);
-    ///    let dtype = array.dtype();
-    ///    assert_eq!(dtype.get_type().name().unwrap().to_string(), "float64");
-    /// });
-    /// ```
-    pub fn get_type(&self) -> &PyType {
-        let dtype_type_ptr = unsafe { *self.as_dtype_ptr() }.typeobj;
-        unsafe { PyType::from_type_ptr(self.py(), dtype_type_ptr) }
-    }
-
     /// Shortcut for creating a descriptor of 'object' type.
     pub fn object(py: Python) -> &Self {
         Self::from_npy_type(py, NPY_TYPES::NPY_OBJECT)
@@ -94,12 +90,207 @@ impl PyArrayDescr {
         }
     }
 
-    /// Retrieves the
-    /// [enumerated type](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types)
-    /// for this type descriptor.
-    pub fn get_typenum(&self) -> c_int {
+    /// Returns the
+    /// [array scalar](https://numpy.org/doc/stable/reference/arrays.scalars.html)
+    /// corresponding to this dtype.
+    ///
+    /// Equivalent to [`np.dtype.type`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.type.html).
+    pub fn typeobj(&self) -> &PyType {
+        let dtype_type_ptr = unsafe { *self.as_dtype_ptr() }.typeobj;
+        unsafe { PyType::from_type_ptr(self.py(), dtype_type_ptr) }
+    }
+
+    /// Returns a unique number for each of the 21 different built-in
+    /// [enumerated types](https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types).
+    ///
+    /// These are roughly ordered from least-to-most precision.
+    ///
+    /// Equivalent to [`np.dtype.num`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.num.html).
+    pub fn num(&self) -> c_int {
         unsafe { *self.as_dtype_ptr() }.type_num
     }
+
+    /// Returns the element size of this data-type object.
+    ///
+    /// Equivalent to [`np.dtype.itemsize`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.itemsize.html).
+    pub fn itemsize(&self) -> usize {
+        unsafe { *self.as_dtype_ptr() }.elsize.max(0) as _
+    }
+
+    /// Returns the required alignment (bytes) of this data-type according to the compiler
+    ///
+    /// Equivalent to [`np.dtype.alignment`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.alignment.html).
+    pub fn alignment(&self) -> usize {
+        unsafe { *self.as_dtype_ptr() }.alignment.max(0) as _
+    }
+
+    /// Returns a character indicating the byte-order of this data-type object.
+    ///
+    /// All built-in data-type objects have byteorder either `=` or `|`.
+    ///
+    /// Equivalent to [`np.dtype.byteorder`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.byteorder.html).
+    pub fn byteorder(&self) -> u8 {
+        unsafe { *self.as_dtype_ptr() }.byteorder.max(0) as _
+    }
+
+    /// Returns a unique character code for each of the 21 different built-in types.
+    ///
+    /// Note: structured data types are categorized as `V` (void).
+    ///
+    /// Equivalent to [`np.dtype.char`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.char.html)
+    pub fn char(&self) -> u8 {
+        unsafe { *self.as_dtype_ptr() }.type_.max(0) as _
+    }
+
+    /// Returns a character code (one of `biufcmMOSUV`) identifying the general kind of data.
+    ///
+    /// Note: structured data types are categorized as `V` (void).
+    ///
+    /// Equivalent to [`np.dtype.kind`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html)
+    pub fn kind(&self) -> u8 {
+        unsafe { *self.as_dtype_ptr() }.kind.max(0) as _
+    }
+
+    /// Returns bit-flags describing how this data type is to be interpreted.
+    ///
+    /// Equivalent to [`np.dtype.flags`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.flags.html)
+    pub fn flags(&self) -> c_char {
+        unsafe { *self.as_dtype_ptr() }.flags
+    }
+
+    /// Returns the number of dimensions if this data type describes a sub-array, and `0` otherwise.
+    ///
+    /// Equivalent to [`np.dtype.ndim`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.ndim.html)
+    pub fn ndim(&self) -> usize {
+        if !self.has_subarray() {
+            return 0;
+        }
+        unsafe { PyTuple_Size((*((*self.as_dtype_ptr()).subarray)).shape).max(0) as _ }
+    }
+
+    /// Returns dtype for the base element of subarrays, regardless of their dimension or shape.
+    ///
+    /// Equivalent to [`np.dtype.base`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.base.html).
+    pub fn base(&self) -> Option<&PyArrayDescr> {
+        if !self.has_subarray() {
+            return None;
+        }
+        Some(unsafe { Self::from_borrowed_ptr(self.py(), (*self.as_dtype_ptr()).subarray as _) })
+    }
+
+    /// Returns shape tuple of the sub-array if this dtype is a sub-array, and `None` otherwise.
+    ///
+    /// Equivalent to [`np.dtype.shape`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.shape.html)
+    pub fn shape(&self) -> Option<Vec<usize>> {
+        if !self.has_subarray() {
+            return None;
+        }
+        Some(
+            // Panic-wise: numpy guarantees that shape is a tuple of non-negative integers
+            unsafe {
+                PyTuple::from_borrowed_ptr(self.py(), (*(*self.as_dtype_ptr()).subarray).shape)
+            }
+            .extract()
+            .unwrap(),
+        )
+    }
+
+    /// Returns `(item_dtype, shape)` if this dtype describes a sub-array, and `None` otherwise.
+    ///
+    /// The `shape` is the fixed shape of the sub-array described by this data type,
+    /// and `item_dtype` the data type of the array.
+    ///
+    /// If a field whose dtype object has this attribute is retrieved, then the extra dimensions
+    /// implied by shape are tacked on to the end of the retrieved array.
+    ///
+    /// Equivalent to [`np.dtype.subdtype`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.subdtype.html)
+    pub fn subdtype(&self) -> Option<(&PyArrayDescr, Vec<usize>)> {
+        self.shape()
+            .and_then(|shape| self.base().map(|base| (base, shape)))
+    }
+
+    /// Returns true if the dtype is a sub-array at the top level.
+    ///
+    /// Equivalent to [`np.dtype.hasobject`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.hasobject.html)
+    pub fn has_object(&self) -> bool {
+        self.flags() & NPY_ITEM_HASOBJECT != 0
+    }
+
+    /// Returns true if the dtype is a struct which maintains field alignment.
+    ///
+    /// This flag is sticky, so when combining multiple structs together, it is preserved
+    /// and produces new dtypes which are also aligned.
+    ///
+    /// Equivalent to [`np.dtype.isalignedstruct`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.isalignedstruct.html)
+    pub fn is_aligned_struct(&self) -> bool {
+        self.flags() & NPY_ALIGNED_STRUCT != 0
+    }
+
+    /// Returns true if the data type is a sub-array.
+    pub fn has_subarray(&self) -> bool {
+        // equivalent to PyDataType_HASSUBARRAY(self)
+        unsafe { !(*self.as_dtype_ptr()).subarray.is_null() }
+    }
+
+    /// Returns true if the data type is a structured type.
+    pub fn has_fields(&self) -> bool {
+        // equivalent to PyDataType_HASFIELDS(self)
+        unsafe { !(*self.as_dtype_ptr()).names.is_null() }
+    }
+
+    /// Returns true if the data type is unsized
+    pub fn is_unsized(&self) -> bool {
+        // equivalent to PyDataType_ISUNSIZED(self)
+        self.itemsize() == 0 && !self.has_fields()
+    }
+
+    /// Returns true if data type byteorder is native, or `None` if not applicable.
+    pub fn is_native_byteorder(&self) -> Option<bool> {
+        // based on PyArray_ISNBO(self->byteorder)
+        match self.byteorder() {
+            b'=' => Some(true),
+            b'|' => None,
+            byteorder if byteorder == NPY_BYTEORDER_CHAR::NPY_NATBYTE as u8 => Some(true),
+            _ => Some(false),
+        }
+    }
+
+    /// Returns an ordered list of field names, or `None` if there are no fields.
+    ///
+    /// The names are ordered according to increasing byte offset.
+    ///
+    /// Equivalent to [`np.dtype.names`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.names.html).
+    pub fn names(&self) -> Option<Vec<&str>> {
+        if !self.has_fields() {
+            return None;
+        }
+        let names = unsafe { PyTuple::from_borrowed_ptr(self.py(), (*self.as_dtype_ptr()).names) };
+        FromPyObject::extract(names).ok()
+    }
+
+    /// Returns names, types and offsets of fields, or `None` if not a structured type.
+    ///
+    /// The iterator has entries in the form `(name, (dtype, offset))` so it can be
+    /// collected directly into a map-like structure.
+    ///
+    /// Note: titles (the optional 3rd tuple element) are ignored.
+    ///
+    /// Equivalent to [`np.dtype.fields`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.fields.html).
+    pub fn fields(&self) -> Option<impl Iterator<Item = (&str, (&PyArrayDescr, usize))> + '_> {
+        if !self.has_fields() {
+            return None;
+        }
+        let dict = unsafe { PyDict::from_borrowed_ptr(self.py(), (*self.as_dtype_ptr()).fields) };
+        // Panic-wise: numpy guarantees that fields are tuples of proper size and type
+        Some(dict.iter().map(|(k, v)| {
+            let name = FromPyObject::extract(k).unwrap();
+            let tuple = v.downcast::<PyTuple>().unwrap();
+            // note: we can't just extract the entire tuple since 3rd element can be a title
+            let dtype = FromPyObject::extract(tuple.as_ref().get_item(0).unwrap()).unwrap();
+            let offset = FromPyObject::extract(tuple.as_ref().get_item(1).unwrap()).unwrap();
+            (name, (dtype, offset))
+        }))
+    }
 }
 
 /// Represents that a type can be an element of `PyArray`.
@@ -259,7 +450,7 @@ mod tests {
     #[test]
     fn test_dtype_names() {
         fn type_name<T: Element>(py: pyo3::Python) -> &str {
-            dtype::<T>(py).get_type().name().unwrap()
+            dtype::<T>(py).typeobj().name().unwrap()
         }
         pyo3::Python::with_gil(|py| {
             assert_eq!(type_name::<bool>(py), "bool_");

src/npyffi/array.rs

@@ -394,6 +394,12 @@ pub unsafe fn PyArray_CheckExact(op: *mut PyObject) -> c_int {
     (ffi::Py_TYPE(op) == PY_ARRAY_API.get_type_object(NpyTypes::PyArray_Type)) as _
 }
 
+// these are under `#if NPY_USE_PYMEM == 1` which seems to be always defined as 1
+pub use pyo3::ffi::{
+    PyMem_RawFree as PyArray_free, PyMem_RawMalloc as PyArray_malloc,
+    PyMem_RawRealloc as PyArray_realloc,
+};
+
 #[cfg(test)]
 mod tests {
     use super::PY_ARRAY_API;

src/npyffi/flags.rs

@@ -1,4 +1,4 @@
-use super::npy_uint32;
+use super::{npy_char, npy_uint32};
 use std::os::raw::c_int;
 
 pub const NPY_ARRAY_C_CONTIGUOUS: c_int = 0x0001;
@@ -62,3 +62,21 @@ pub const NPY_ITER_OVERLAP_ASSUME_ELEMENTWISE: npy_uint32 = 0x40000000;
 
 pub const NPY_ITER_GLOBAL_FLAGS: npy_uint32 = 0x0000ffff;
 pub const NPY_ITER_PER_OP_FLAGS: npy_uint32 = 0xffff0000;
+
+pub const NPY_ITEM_REFCOUNT: npy_char = 0x01;
+pub const NPY_ITEM_HASOBJECT: npy_char = 0x01;
+pub const NPY_LIST_PICKLE: npy_char = 0x02;
+pub const NPY_ITEM_IS_POINTER: npy_char = 0x04;
+pub const NPY_NEEDS_INIT: npy_char = 0x08;
+pub const NPY_NEEDS_PYAPI: npy_char = 0x10;
+pub const NPY_USE_GETITEM: npy_char = 0x20;
+pub const NPY_USE_SETITEM: npy_char = 0x40;
+pub const NPY_ALIGNED_STRUCT: npy_char = -128; // 0x80
+pub const NPY_FROM_FIELDS: npy_char =
+    NPY_NEEDS_INIT | NPY_LIST_PICKLE | NPY_ITEM_REFCOUNT | NPY_NEEDS_PYAPI;
+pub const NPY_OBJECT_DTYPE_FLAGS: npy_char = NPY_LIST_PICKLE
+    | NPY_USE_GETITEM
+    | NPY_ITEM_IS_POINTER
+    | NPY_ITEM_REFCOUNT
+    | NPY_NEEDS_INIT
+    | NPY_NEEDS_PYAPI;

src/npyffi/objects.rs

@@ -398,3 +398,17 @@ pub struct NpyAuxData {
 
 pub type NpyAuxData_FreeFunc = Option<unsafe extern "C" fn(*mut NpyAuxData)>;
 pub type NpyAuxData_CloneFunc = Option<unsafe extern "C" fn(*mut NpyAuxData) -> *mut NpyAuxData>;
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct PyArray_DatetimeMetaData {
+    pub base: NPY_DATETIMEUNIT,
+    pub num: c_int,
+}
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct PyArray_DatetimeDTypeMetaData {
+    pub base: NpyAuxData,
+    pub meta: PyArray_DatetimeMetaData,
+}