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,22 @@
+use std::collections::BTreeMap;
 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, PyResult,
+};
 
-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};
 
@@ -94,12 +106,202 @@ 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 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(
+            // TODO: can this be done simpler, without the incref?
+            unsafe {
+                PyTuple::from_borrowed_ptr(self.py(), (*(*self.as_dtype_ptr()).subarray).shape)
+            }
+            .extract()
+            .unwrap(), // TODO: unwrap? numpy sort-of guarantees it will be an int tuple
+        )
+    }
+
+    /// 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) };
+        <_>::extract(names).ok()
+    }
+
+    /// Returns a dictionary of fields, or `None` if not a structured type.
+    ///
+    /// The dictionary is indexed by keys that are the names of the fields. Each entry in
+    /// the dictionary is a tuple fully describing the field: `(dtype, offset)`.
+    ///
+    /// 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<BTreeMap<&str, (&PyArrayDescr, usize)>> {
+        if !self.has_fields() {
+            return None;
+        }
+        // TODO: can this be done simpler, without the incref?
+        let dict = unsafe { PyDict::from_borrowed_ptr(self.py(), (*self.as_dtype_ptr()).fields) };
+        let mut fields = BTreeMap::new();
+        (|| -> PyResult<_> {
+            for (k, v) in dict.iter() {
+                // TODO: alternatively, could unwrap everything here
+                let name = <_>::extract(k)?;
+                let tuple = v.downcast::<PyTuple>()?;
+                let dtype = <_>::extract(tuple.as_ref().get_item(0)?)?;
+                let offset = <_>::extract(tuple.as_ref().get_item(1)?)?;
+                fields.insert(name, (dtype, offset));
+            }
+            Ok(fields)
+        })()
+        .ok()
+    }
 }
 
 /// Represents that a type can be an element of `PyArray`.

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,
+}

src/npyffi/types.rs

@@ -115,7 +115,7 @@ pub enum NPY_DATETIMEUNIT {
 }
 
 #[repr(u32)]
-#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub enum NPY_TYPES {
     NPY_BOOL = 0,
     NPY_BYTE = 1,
@@ -201,3 +201,70 @@ pub struct npy_stride_sort_item {
     pub perm: npy_intp,
     pub stride: npy_intp,
 }
+
+#[repr(u8)]
+#[derive(Debug, Clone, Copy)]
+pub enum NPY_TYPECHAR {
+    NPY_BOOLLTR = b'?',
+    NPY_BYTELTR = b'b',
+    NPY_UBYTELTR = b'B',
+    NPY_SHORTLTR = b'h',
+    NPY_USHORTLTR = b'H',
+    NPY_INTLTR = b'i',
+    NPY_UINTLTR = b'I',
+    NPY_LONGLTR = b'l',
+    NPY_ULONGLTR = b'L',
+    NPY_LONGLONGLTR = b'q',
+    NPY_ULONGLONGLTR = b'Q',
+    NPY_HALFLTR = b'e',
+    NPY_FLOATLTR = b'f',
+    NPY_DOUBLELTR = b'd',
+    NPY_LONGDOUBLELTR = b'g',
+    NPY_CFLOATLTR = b'F',
+    NPY_CDOUBLELTR = b'D',
+    NPY_CLONGDOUBLELTR = b'G',
+    NPY_OBJECTLTR = b'O',
+    NPY_STRINGLTR = b'S',
+    NPY_STRINGLTR2 = b'a',
+    NPY_UNICODELTR = b'U',
+    NPY_VOIDLTR = b'V',
+    NPY_DATETIMELTR = b'M',
+    NPY_TIMEDELTALTR = b'm',
+    NPY_CHARLTR = b'c',
+    NPY_INTPLTR = b'p',
+    NPY_UINTPLTR = b'P',
+}
+
+#[repr(u8)]
+#[derive(Debug, Clone, Copy)]
+// NPY_TYPEKINDCHAR doesn't exist in the header, but these enum values are not
+// related to NPY_TYPECHAR although being stuffed into it (type kinds, not type codes)
+pub enum NPY_TYPEKINDCHAR {
+    NPY_GENBOOLLTR = b'b',
+    NPY_SIGNEDLTR = b'i',
+    NPY_UNSIGNEDLTR = b'u',
+    NPY_FLOATINGLTR = b'f',
+    NPY_COMPLEXLTR = b'c',
+}
+
+#[repr(u8)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub enum NPY_BYTEORDER_CHAR {
+    NPY_LITTLE = b'<',
+    NPY_BIG = b'>',
+    NPY_NATIVE = b'=',
+    NPY_SWAP = b's',
+    NPY_IGNORE = b'|',
+}
+
+impl NPY_BYTEORDER_CHAR {
+    #[cfg(target_endian = "little")]
+    pub const NPY_NATBYTE: Self = Self::NPY_LITTLE;
+    #[cfg(target_endian = "little")]
+    pub const NPY_OPPBYTE: Self = Self::NPY_BIG;
+
+    #[cfg(target_endian = "big")]
+    pub const NPY_NATBYTE: Self = Self::NPY_BIG;
+    #[cfg(target_endian = "big")]
+    pub const NPY_OPPBYTE: Self = Self::NPY_LITTLE;
+}