scirs2_numpy/
dtype.rs

1use std::mem::size_of;
2use std::os::raw::{c_int, c_long, c_longlong, c_short, c_uint, c_ulong, c_ulonglong, c_ushort};
3use std::ptr;
4
5#[cfg(feature = "half")]
6use half::{bf16, f16};
7use num_traits::{Bounded, Zero};
8#[cfg(feature = "half")]
9use pyo3::sync::PyOnceLock;
10use pyo3::{
11    conversion::IntoPyObject,
12    exceptions::{PyIndexError, PyValueError},
13    ffi::{self, PyTuple_Size},
14    pyobject_native_type_named,
15    types::{PyAnyMethods, PyDict, PyDictMethods, PyTuple, PyType},
16    Borrowed, Bound, Py, PyAny, PyResult, PyTypeInfo, Python,
17};
18
19use crate::npyffi::{
20    NpyTypes, PyArray_Descr, PyDataType_ALIGNMENT, PyDataType_ELSIZE, PyDataType_FIELDS,
21    PyDataType_FLAGS, PyDataType_NAMES, PyDataType_SUBARRAY, NPY_ALIGNED_STRUCT,
22    NPY_BYTEORDER_CHAR, NPY_ITEM_HASOBJECT, NPY_TYPES, PY_ARRAY_API,
23};
24
25pub use num_complex::{Complex32, Complex64};
26
27/// Binding of [`numpy.dtype`][dtype].
28///
29/// # Example
30///
31/// ```
32/// use numpy::{dtype, get_array_module, PyArrayDescr, PyArrayDescrMethods};
33/// use numpy::pyo3::{types::{IntoPyDict, PyAnyMethods}, Python, ffi::c_str};
34///
35/// # fn main() -> pyo3::PyResult<()> {
36/// Python::attach(|py| {
37///     let locals = [("np", get_array_module(py)?)].into_py_dict(py)?;
38///
39///     let dt = py
40///         .eval(c_str!("np.array([1, 2, 3.0]).dtype"), Some(&locals), None)?
41///         .cast_into::<PyArrayDescr>()?;
42///
43///     assert!(dt.is_equiv_to(&dtype::<f64>(py)));
44/// #   Ok(())
45/// })
46/// # }
47/// ```
48///
49/// [dtype]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.html
50#[repr(transparent)]
51pub struct PyArrayDescr(PyAny);
52
53pyobject_native_type_named!(PyArrayDescr);
54
55unsafe impl PyTypeInfo for PyArrayDescr {
56    const NAME: &'static str = "PyArrayDescr";
57    const MODULE: Option<&'static str> = Some("numpy");
58
59    #[inline]
60    fn type_object_raw<'py>(py: Python<'py>) -> *mut ffi::PyTypeObject {
61        unsafe { PY_ARRAY_API.get_type_object(py, NpyTypes::PyArrayDescr_Type) }
62    }
63}
64
65/// Returns the type descriptor ("dtype") for a registered type.
66#[inline]
67pub fn dtype<'py, T: Element>(py: Python<'py>) -> Bound<'py, PyArrayDescr> {
68    T::get_dtype(py)
69}
70
71impl PyArrayDescr {
72    /// Creates a new type descriptor ("dtype") object from an arbitrary object.
73    ///
74    /// Equivalent to invoking the constructor of [`numpy.dtype`][dtype].
75    ///
76    /// [dtype]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.html
77    #[inline]
78    pub fn new<'a, 'py, T>(py: Python<'py>, ob: T) -> PyResult<Bound<'py, Self>>
79    where
80        T: IntoPyObject<'py>,
81    {
82        fn inner<'py>(
83            py: Python<'py>,
84            obj: Borrowed<'_, 'py, PyAny>,
85        ) -> PyResult<Bound<'py, PyArrayDescr>> {
86            let mut descr: *mut PyArray_Descr = ptr::null_mut();
87            unsafe {
88                // None is an invalid input here and is not converted to NPY_DEFAULT_TYPE
89                PY_ARRAY_API.PyArray_DescrConverter2(py, obj.as_ptr(), &mut descr);
90                Bound::from_owned_ptr_or_err(py, descr.cast()).map(|any| any.cast_into_unchecked())
91            }
92        }
93
94        inner(
95            py,
96            ob.into_pyobject(py)
97                .map_err(Into::into)?
98                .into_any()
99                .as_borrowed(),
100        )
101    }
102
103    /// Shortcut for creating a type descriptor of `object` type.
104    #[inline]
105    pub fn object(py: Python<'_>) -> Bound<'_, Self> {
106        Self::from_npy_type(py, NPY_TYPES::NPY_OBJECT)
107    }
108
109    /// Returns the type descriptor for a registered type.
110    #[inline]
111    pub fn of<'py, T: Element>(py: Python<'py>) -> Bound<'py, Self> {
112        T::get_dtype(py)
113    }
114
115    fn from_npy_type<'py>(py: Python<'py>, npy_type: NPY_TYPES) -> Bound<'py, Self> {
116        unsafe {
117            let descr = PY_ARRAY_API.PyArray_DescrFromType(py, npy_type as _);
118            Bound::from_owned_ptr(py, descr.cast()).cast_into_unchecked()
119        }
120    }
121
122    pub(crate) fn new_from_npy_type<'py>(py: Python<'py>, npy_type: NPY_TYPES) -> Bound<'py, Self> {
123        unsafe {
124            let descr = PY_ARRAY_API.PyArray_DescrNewFromType(py, npy_type as _);
125            Bound::from_owned_ptr(py, descr.cast()).cast_into_unchecked()
126        }
127    }
128}
129
130/// Implementation of functionality for [`PyArrayDescr`].
131#[doc(alias = "PyArrayDescr")]
132pub trait PyArrayDescrMethods<'py>: Sealed {
133    /// Returns `self` as `*mut PyArray_Descr`.
134    fn as_dtype_ptr(&self) -> *mut PyArray_Descr;
135
136    /// Returns `self` as `*mut PyArray_Descr` while increasing the reference count.
137    ///
138    /// Useful in cases where the descriptor is stolen by the API.
139    fn into_dtype_ptr(self) -> *mut PyArray_Descr;
140
141    /// Returns true if two type descriptors are equivalent.
142    fn is_equiv_to(&self, other: &Self) -> bool;
143
144    /// Returns the [array scalar][arrays-scalars] corresponding to this type descriptor.
145    ///
146    /// Equivalent to [`numpy.dtype.type`][dtype-type].
147    ///
148    /// [arrays-scalars]: https://numpy.org/doc/stable/reference/arrays.scalars.html
149    /// [dtype-type]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.type.html
150    fn typeobj(&self) -> Bound<'py, PyType>;
151
152    /// Returns a unique number for each of the 21 different built-in
153    /// [enumerated types][enumerated-types].
154    ///
155    /// These are roughly ordered from least-to-most precision.
156    ///
157    /// Equivalent to [`numpy.dtype.num`][dtype-num].
158    ///
159    /// [enumerated-types]: https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types
160    /// [dtype-num]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.num.html
161    fn num(&self) -> c_int {
162        unsafe { &*self.as_dtype_ptr() }.type_num
163    }
164
165    /// Returns the element size of this type descriptor.
166    ///
167    /// Equivalent to [`numpy.dtype.itemsize`][dtype-itemsize].
168    ///
169    /// [dtype-itemsiize]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.itemsize.html
170    fn itemsize(&self) -> usize;
171
172    /// Returns the required alignment (bytes) of this type descriptor according to the compiler.
173    ///
174    /// Equivalent to [`numpy.dtype.alignment`][dtype-alignment].
175    ///
176    /// [dtype-alignment]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.alignment.html
177    fn alignment(&self) -> usize;
178
179    /// Returns an ASCII character indicating the byte-order of this type descriptor object.
180    ///
181    /// All built-in data-type objects have byteorder either `=` or `|`.
182    ///
183    /// Equivalent to [`numpy.dtype.byteorder`][dtype-byteorder].
184    ///
185    /// [dtype-byteorder]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.byteorder.html
186    fn byteorder(&self) -> u8 {
187        unsafe { &*self.as_dtype_ptr() }.byteorder.max(0) as _
188    }
189
190    /// Returns a unique ASCII character for each of the 21 different built-in types.
191    ///
192    /// Note that structured data types are categorized as `V` (void).
193    ///
194    /// Equivalent to [`numpy.dtype.char`][dtype-char].
195    ///
196    /// [dtype-char]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.char.html
197    fn char(&self) -> u8 {
198        unsafe { &*self.as_dtype_ptr() }.type_.max(0) as _
199    }
200
201    /// Returns an ASCII character (one of `biufcmMOSUV`) identifying the general kind of data.
202    ///
203    /// Note that structured data types are categorized as `V` (void).
204    ///
205    /// Equivalent to [`numpy.dtype.kind`][dtype-kind].
206    ///
207    /// [dtype-kind]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html
208    fn kind(&self) -> u8 {
209        unsafe { &*self.as_dtype_ptr() }.kind.max(0) as _
210    }
211
212    /// Returns bit-flags describing how this type descriptor is to be interpreted.
213    ///
214    /// Equivalent to [`numpy.dtype.flags`][dtype-flags].
215    ///
216    /// [dtype-flags]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.flags.html
217    fn flags(&self) -> u64;
218
219    /// Returns the number of dimensions if this type descriptor represents a sub-array, and zero otherwise.
220    ///
221    /// Equivalent to [`numpy.dtype.ndim`][dtype-ndim].
222    ///
223    /// [dtype-ndim]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.ndim.html
224    fn ndim(&self) -> usize;
225
226    /// Returns the type descriptor for the base element of subarrays, regardless of their dimension or shape.
227    ///
228    /// If the dtype is not a subarray, returns self.
229    ///
230    /// Equivalent to [`numpy.dtype.base`][dtype-base].
231    ///
232    /// [dtype-base]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.base.html
233    fn base(&self) -> Bound<'py, PyArrayDescr>;
234
235    /// Returns the shape of the sub-array.
236    ///
237    /// If the dtype is not a sub-array, an empty vector is returned.
238    ///
239    /// Equivalent to [`numpy.dtype.shape`][dtype-shape].
240    ///
241    /// [dtype-shape]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.shape.html
242    fn shape(&self) -> Vec<usize>;
243
244    /// Returns true if the type descriptor contains any reference-counted objects in any fields or sub-dtypes.
245    ///
246    /// Equivalent to [`numpy.dtype.hasobject`][dtype-hasobject].
247    ///
248    /// [dtype-hasobject]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.hasobject.html
249    fn has_object(&self) -> bool {
250        self.flags() & NPY_ITEM_HASOBJECT != 0
251    }
252
253    /// Returns true if the type descriptor is a struct which maintains field alignment.
254    ///
255    /// This flag is sticky, so when combining multiple structs together, it is preserved
256    /// and produces new dtypes which are also aligned.
257    ///
258    /// Equivalent to [`numpy.dtype.isalignedstruct`][dtype-isalignedstruct].
259    ///
260    /// [dtype-isalignedstruct]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.isalignedstruct.html
261    fn is_aligned_struct(&self) -> bool {
262        self.flags() & NPY_ALIGNED_STRUCT != 0
263    }
264
265    /// Returns true if the type descriptor is a sub-array.
266    ///
267    /// Equivalent to PyDataType_HASSUBARRAY(self).
268    fn has_subarray(&self) -> bool;
269
270    /// Returns true if the type descriptor is a structured type.
271    ///
272    /// Equivalent to PyDataType_HASFIELDS(self).
273    fn has_fields(&self) -> bool;
274
275    /// Returns true if type descriptor byteorder is native, or `None` if not applicable.
276    fn is_native_byteorder(&self) -> Option<bool> {
277        // based on PyArray_ISNBO(self->byteorder)
278        match self.byteorder() {
279            b'=' => Some(true),
280            b'|' => None,
281            byteorder => Some(byteorder == NPY_BYTEORDER_CHAR::NPY_NATBYTE as u8),
282        }
283    }
284
285    /// Returns an ordered list of field names, or `None` if there are no fields.
286    ///
287    /// The names are ordered according to increasing byte offset.
288    ///
289    /// Equivalent to [`numpy.dtype.names`][dtype-names].
290    ///
291    /// [dtype-names]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.names.html
292    fn names(&self) -> Option<Vec<String>>;
293
294    /// Returns the type descriptor and offset of the field with the given name.
295    ///
296    /// This method will return an error if this type descriptor is not structured,
297    /// or if it does not contain a field with a given name.
298    ///
299    /// The list of all names can be found via `PyArrayDescr::names`.
300    ///
301    /// Equivalent to retrieving a single item from [`numpy.dtype.fields`][dtype-fields].
302    ///
303    /// [dtype-fields]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.fields.html
304    fn get_field(&self, name: &str) -> PyResult<(Bound<'py, PyArrayDescr>, usize)>;
305}
306
307mod sealed {
308    pub trait Sealed {}
309}
310
311use sealed::Sealed;
312
313impl<'py> PyArrayDescrMethods<'py> for Bound<'py, PyArrayDescr> {
314    fn as_dtype_ptr(&self) -> *mut PyArray_Descr {
315        self.as_ptr() as _
316    }
317
318    fn into_dtype_ptr(self) -> *mut PyArray_Descr {
319        self.into_ptr() as _
320    }
321
322    fn is_equiv_to(&self, other: &Self) -> bool {
323        let self_ptr = self.as_dtype_ptr();
324        let other_ptr = other.as_dtype_ptr();
325
326        unsafe {
327            self_ptr == other_ptr
328                || PY_ARRAY_API.PyArray_EquivTypes(self.py(), self_ptr, other_ptr) != 0
329        }
330    }
331
332    fn typeobj(&self) -> Bound<'py, PyType> {
333        let dtype_type_ptr = unsafe { &*self.as_dtype_ptr() }.typeobj;
334        unsafe { PyType::from_borrowed_type_ptr(self.py(), dtype_type_ptr) }
335    }
336
337    fn itemsize(&self) -> usize {
338        unsafe { PyDataType_ELSIZE(self.py(), self.as_dtype_ptr()).max(0) as _ }
339    }
340
341    fn alignment(&self) -> usize {
342        unsafe { PyDataType_ALIGNMENT(self.py(), self.as_dtype_ptr()).max(0) as _ }
343    }
344
345    fn flags(&self) -> u64 {
346        unsafe { PyDataType_FLAGS(self.py(), self.as_dtype_ptr()) as _ }
347    }
348
349    fn ndim(&self) -> usize {
350        let subarray = unsafe { PyDataType_SUBARRAY(self.py(), self.as_dtype_ptr()).as_ref() };
351        match subarray {
352            None => 0,
353            Some(subarray) => unsafe { PyTuple_Size(subarray.shape) }.max(0) as _,
354        }
355    }
356
357    fn base(&self) -> Bound<'py, PyArrayDescr> {
358        let subarray = unsafe { PyDataType_SUBARRAY(self.py(), self.as_dtype_ptr()).as_ref() };
359        match subarray {
360            None => self.clone(),
361            Some(subarray) => unsafe {
362                Bound::from_borrowed_ptr(self.py(), subarray.base.cast()).cast_into_unchecked()
363            },
364        }
365    }
366
367    fn shape(&self) -> Vec<usize> {
368        let subarray = unsafe { PyDataType_SUBARRAY(self.py(), self.as_dtype_ptr()).as_ref() };
369        match subarray {
370            None => Vec::new(),
371            Some(subarray) => {
372                // NumPy guarantees that shape is a tuple of non-negative integers so this should never panic.
373                let shape = unsafe { Borrowed::from_ptr(self.py(), subarray.shape) };
374                shape.extract().expect("Operation failed")
375            }
376        }
377    }
378
379    fn has_subarray(&self) -> bool {
380        unsafe { !PyDataType_SUBARRAY(self.py(), self.as_dtype_ptr()).is_null() }
381    }
382
383    fn has_fields(&self) -> bool {
384        unsafe { !PyDataType_NAMES(self.py(), self.as_dtype_ptr()).is_null() }
385    }
386
387    fn names(&self) -> Option<Vec<String>> {
388        if !self.has_fields() {
389            return None;
390        }
391        let names = unsafe {
392            Borrowed::from_ptr(self.py(), PyDataType_NAMES(self.py(), self.as_dtype_ptr()))
393        };
394        names.extract().ok()
395    }
396
397    fn get_field(&self, name: &str) -> PyResult<(Bound<'py, PyArrayDescr>, usize)> {
398        if !self.has_fields() {
399            return Err(PyValueError::new_err(
400                "cannot get field information: type descriptor has no fields",
401            ));
402        }
403        let dict = unsafe {
404            Borrowed::from_ptr(self.py(), PyDataType_FIELDS(self.py(), self.as_dtype_ptr()))
405        };
406        let dict = unsafe { dict.cast_unchecked::<PyDict>() };
407        // NumPy guarantees that fields are tuples of proper size and type, so this should never panic.
408        let tuple = dict
409            .get_item(name)?
410            .ok_or_else(|| PyIndexError::new_err(name.to_owned()))?
411            .cast_into::<PyTuple>()
412            .expect("Operation failed");
413        // Note that we cannot just extract the entire tuple since the third element can be a title.
414        let dtype = tuple
415            .get_item(0)
416            .expect("Operation failed")
417            .cast_into::<PyArrayDescr>()
418            .expect("Operation failed");
419        let offset = tuple
420            .get_item(1)
421            .expect("Operation failed")
422            .extract()
423            .expect("Operation failed");
424        Ok((dtype, offset))
425    }
426}
427
428impl Sealed for Bound<'_, PyArrayDescr> {}
429
430/// Represents that a type can be an element of `PyArray`.
431///
432/// Currently, only integer/float/complex/object types are supported. The [NumPy documentation][enumerated-types]
433/// list the other built-in types which we are not yet implemented.
434///
435/// Note that NumPy's integer types like `numpy.int_` and `numpy.uint` are based on C's integer hierarchy
436/// which implies that their widths change depending on the platform's [data model][data-models].
437/// For example, `numpy.int_` matches C's `long` which is 32 bits wide on Windows (using the LLP64 data model)
438/// but 64 bits wide on Linux (using the LP64 data model).
439///
440/// In contrast, Rust's [`isize`] and [`usize`] types are defined to have the same width as a pointer
441/// and are therefore always 64 bits wide on 64-bit platforms. If you want to match NumPy's behaviour,
442/// consider using the [`c_long`][std::ffi::c_long] and [`c_ulong`][std::ffi::c_ulong] type aliases.
443///
444/// # Safety
445///
446/// A type `T` that implements this trait should be safe when managed by a NumPy
447/// array, thus implementing this trait is marked unsafe. Data types that don't
448/// contain Python objects (i.e., either the object type itself or record types
449/// containing object-type fields) are assumed to be trivially copyable, which
450/// is reflected in the `IS_COPY` flag. Furthermore, it is assumed that for
451/// the object type the elements are pointers into the Python heap and that the
452/// corresponding `Clone` implemenation will never panic as it only increases
453/// the reference count.
454///
455/// # Custom element types
456///
457/// Note that we cannot safely store `Py<T>` where `T: PyClass`, because the type information would be
458/// eliminated in the resulting NumPy array.
459/// In other words, objects are always treated as `Py<PyAny>` (a.k.a. `PyObject`) by Python code,
460/// and only `Py<PyAny>` can be stored in a type safe manner.
461///
462/// You can however create [`Array<Py<T>, D>`][ndarray::Array] and turn that into a NumPy array
463/// safely and efficiently using [`from_owned_object_array`][crate::PyArray::from_owned_object_array].
464///
465/// [enumerated-types]: https://numpy.org/doc/stable/reference/c-api/dtype.html#enumerated-types
466/// [data-models]: https://en.wikipedia.org/wiki/64-bit_computing#64-bit_data_models
467pub unsafe trait Element: Sized + Send + Sync {
468    /// Flag that indicates whether this type is trivially copyable.
469    ///
470    /// It should be set to true for all trivially copyable types (like scalar types
471    /// and record/array types only containing trivially copyable fields and elements).
472    ///
473    /// This flag should *always* be set to `false` for object types or record types
474    /// that contain object-type fields.
475    const IS_COPY: bool;
476
477    /// Returns the associated type descriptor ("dtype") for the given element type.
478    fn get_dtype(py: Python<'_>) -> Bound<'_, PyArrayDescr>;
479
480    /// Create a clone of the value while the GIL is guaranteed to be held.
481    fn clone_ref(&self, py: Python<'_>) -> Self;
482
483    /// Create an owned copy of the slice while the GIL is guaranteed to be held.
484    ///
485    /// Some types may provide implementations of this method that are more efficient
486    /// than simply mapping the `py_clone` method to each element in the slice.
487    #[inline]
488    fn vec_from_slice(py: Python<'_>, slc: &[Self]) -> Vec<Self> {
489        slc.iter().map(|elem| elem.clone_ref(py)).collect()
490    }
491
492    /// Create an owned copy of the array while the GIL is guaranteed to be held.
493    ///
494    /// Some types may provide implementations of this method that are more efficient
495    /// than simply mapping the `py_clone` method to each element in the view.
496    #[inline]
497    fn array_from_view<D>(
498        py: Python<'_>,
499        view: ::ndarray::ArrayView<'_, Self, D>,
500    ) -> ::ndarray::Array<Self, D>
501    where
502        D: ::ndarray::Dimension,
503    {
504        view.map(|elem| elem.clone_ref(py))
505    }
506}
507
508fn npy_int_type_lookup<T, T0, T1, T2>(npy_types: [NPY_TYPES; 3]) -> NPY_TYPES {
509    // `npy_common.h` defines the integer aliases. In order, it checks:
510    // NPY_BITSOF_LONG, NPY_BITSOF_LONGLONG, NPY_BITSOF_INT, NPY_BITSOF_SHORT, NPY_BITSOF_CHAR
511    // and assigns the alias to the first matching size, so we should check in this order.
512    match size_of::<T>() {
513        x if x == size_of::<T0>() => npy_types[0],
514        x if x == size_of::<T1>() => npy_types[1],
515        x if x == size_of::<T2>() => npy_types[2],
516        _ => panic!("Unable to match integer type descriptor: {npy_types:?}"),
517    }
518}
519
520fn npy_int_type<T: Bounded + Zero + Sized + PartialEq>() -> NPY_TYPES {
521    let is_unsigned = T::min_value() == T::zero();
522    let bit_width = 8 * size_of::<T>();
523
524    match (is_unsigned, bit_width) {
525        (false, 8) => NPY_TYPES::NPY_BYTE,
526        (false, 16) => NPY_TYPES::NPY_SHORT,
527        (false, 32) => npy_int_type_lookup::<i32, c_long, c_int, c_short>([
528            NPY_TYPES::NPY_LONG,
529            NPY_TYPES::NPY_INT,
530            NPY_TYPES::NPY_SHORT,
531        ]),
532        (false, 64) => npy_int_type_lookup::<i64, c_long, c_longlong, c_int>([
533            NPY_TYPES::NPY_LONG,
534            NPY_TYPES::NPY_LONGLONG,
535            NPY_TYPES::NPY_INT,
536        ]),
537        (true, 8) => NPY_TYPES::NPY_UBYTE,
538        (true, 16) => NPY_TYPES::NPY_USHORT,
539        (true, 32) => npy_int_type_lookup::<u32, c_ulong, c_uint, c_ushort>([
540            NPY_TYPES::NPY_ULONG,
541            NPY_TYPES::NPY_UINT,
542            NPY_TYPES::NPY_USHORT,
543        ]),
544        (true, 64) => npy_int_type_lookup::<u64, c_ulong, c_ulonglong, c_uint>([
545            NPY_TYPES::NPY_ULONG,
546            NPY_TYPES::NPY_ULONGLONG,
547            NPY_TYPES::NPY_UINT,
548        ]),
549        _ => unreachable!(),
550    }
551}
552
553// Invoke within the `Element` impl for a `Clone` type to provide an efficient
554// implementation of the cloning methods
555macro_rules! clone_methods_impl {
556    ($Self:ty) => {
557        #[inline]
558        fn clone_ref(&self, _py: ::pyo3::Python<'_>) -> $Self {
559            ::std::clone::Clone::clone(self)
560        }
561
562        #[inline]
563        fn vec_from_slice(_py: ::pyo3::Python<'_>, slc: &[$Self]) -> Vec<$Self> {
564            ::std::borrow::ToOwned::to_owned(slc)
565        }
566
567        #[inline]
568        fn array_from_view<D>(
569            _py: ::pyo3::Python<'_>,
570            view: ::ndarray::ArrayView<'_, $Self, D>,
571        ) -> ::ndarray::Array<$Self, D>
572        where
573            D: ::ndarray::Dimension,
574        {
575            ::ndarray::ArrayView::to_owned(&view)
576        }
577    };
578}
579pub(crate) use clone_methods_impl;
580use pyo3::BoundObject;
581
582macro_rules! impl_element_scalar {
583    (@impl: $ty:ty, $npy_type:expr $(,#[$meta:meta])*) => {
584        $(#[$meta])*
585        unsafe impl Element for $ty {
586            const IS_COPY: bool = true;
587
588            fn get_dtype(py: Python<'_>) -> Bound<'_, PyArrayDescr> {
589                PyArrayDescr::from_npy_type(py, $npy_type)
590            }
591
592            clone_methods_impl!($ty);
593        }
594    };
595    ($ty:ty => $npy_type:ident $(,#[$meta:meta])*) => {
596        impl_element_scalar!(@impl: $ty, NPY_TYPES::$npy_type $(,#[$meta])*);
597    };
598    ($($tys:ty),+) => {
599        $(impl_element_scalar!(@impl: $tys, npy_int_type::<$tys>());)+
600    };
601}
602
603impl_element_scalar!(bool => NPY_BOOL);
604
605impl_element_scalar!(i8, i16, i32, i64);
606impl_element_scalar!(u8, u16, u32, u64);
607
608impl_element_scalar!(f32 => NPY_FLOAT);
609impl_element_scalar!(f64 => NPY_DOUBLE);
610
611#[cfg(feature = "half")]
612impl_element_scalar!(f16 => NPY_HALF);
613
614#[cfg(feature = "half")]
615unsafe impl Element for bf16 {
616    const IS_COPY: bool = true;
617
618    fn get_dtype(py: Python<'_>) -> Bound<'_, PyArrayDescr> {
619        static DTYPE: PyOnceLock<Py<PyArrayDescr>> = PyOnceLock::new();
620
621        DTYPE
622            .get_or_init(py, || {
623                PyArrayDescr::new(py, "bfloat16").expect("A package which provides a `bfloat16` data type for NumPy is required to use the `half::bf16` element type.").unbind()
624            })
625            .clone_ref(py)
626            .into_bound(py)
627    }
628
629    clone_methods_impl!(Self);
630}
631
632impl_element_scalar!(Complex32 => NPY_CFLOAT,
633    #[doc = "Complex type with `f32` components which maps to `numpy.csingle` (`numpy.complex64`)."]);
634impl_element_scalar!(Complex64 => NPY_CDOUBLE,
635    #[doc = "Complex type with `f64` components which maps to `numpy.cdouble` (`numpy.complex128`)."]);
636
637#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
638impl_element_scalar!(usize, isize);
639
640unsafe impl Element for Py<PyAny> {
641    const IS_COPY: bool = false;
642
643    fn get_dtype(py: Python<'_>) -> Bound<'_, PyArrayDescr> {
644        PyArrayDescr::object(py)
645    }
646
647    #[inline]
648    fn clone_ref(&self, py: Python<'_>) -> Self {
649        Py::clone_ref(self, py)
650    }
651}
652
653#[cfg(test)]
654mod tests {
655    use super::*;
656
657    use pyo3::types::PyString;
658    use pyo3::{py_run, types::PyTypeMethods};
659
660    use crate::npyffi::{is_numpy_2, NPY_NEEDS_PYAPI};
661
662    #[test]
663    fn test_dtype_new() {
664        Python::attach(|py| {
665            assert!(PyArrayDescr::new(py, "float64")
666                .expect("Operation failed")
667                .is(dtype::<f64>(py)));
668
669            let dt =
670                PyArrayDescr::new(py, [("a", "O"), ("b", "?")].as_ref()).expect("Operation failed");
671            assert_eq!(dt.names(), Some(vec!["a".to_owned(), "b".to_owned()]));
672            assert!(dt.has_object());
673            assert!(dt
674                .get_field("a")
675                .expect("Operation failed")
676                .0
677                .is(dtype::<Py<PyAny>>(py)));
678            assert!(dt
679                .get_field("b")
680                .expect("Operation failed")
681                .0
682                .is(dtype::<bool>(py)));
683
684            assert!(PyArrayDescr::new(py, 123_usize).is_err());
685        });
686    }
687
688    #[test]
689    fn test_dtype_names() {
690        fn type_name<T: Element>(py: Python<'_>) -> Bound<'_, PyString> {
691            dtype::<T>(py)
692                .typeobj()
693                .qualname()
694                .expect("Operation failed")
695        }
696        Python::attach(|py| {
697            if is_numpy_2(py) {
698                assert_eq!(type_name::<bool>(py), "bool");
699            } else {
700                assert_eq!(type_name::<bool>(py), "bool_");
701            }
702
703            assert_eq!(type_name::<i8>(py), "int8");
704            assert_eq!(type_name::<i16>(py), "int16");
705            assert_eq!(type_name::<i32>(py), "int32");
706            assert_eq!(type_name::<i64>(py), "int64");
707            assert_eq!(type_name::<u8>(py), "uint8");
708            assert_eq!(type_name::<u16>(py), "uint16");
709            assert_eq!(type_name::<u32>(py), "uint32");
710            assert_eq!(type_name::<u64>(py), "uint64");
711            assert_eq!(type_name::<f32>(py), "float32");
712            assert_eq!(type_name::<f64>(py), "float64");
713
714            assert_eq!(type_name::<Complex32>(py), "complex64");
715            assert_eq!(type_name::<Complex64>(py), "complex128");
716
717            #[cfg(target_pointer_width = "32")]
718            {
719                assert_eq!(type_name::<usize>(py), "uint32");
720                assert_eq!(type_name::<isize>(py), "int32");
721            }
722
723            #[cfg(target_pointer_width = "64")]
724            {
725                assert_eq!(type_name::<usize>(py), "uint64");
726                assert_eq!(type_name::<isize>(py), "int64");
727            }
728        });
729    }
730
731    #[test]
732    fn test_dtype_methods_scalar() {
733        Python::attach(|py| {
734            let dt = dtype::<f64>(py);
735
736            assert_eq!(dt.num(), NPY_TYPES::NPY_DOUBLE as c_int);
737            assert_eq!(dt.flags(), 0);
738            assert_eq!(
739                dt.typeobj().qualname().expect("Operation failed"),
740                "float64"
741            );
742            assert_eq!(dt.char(), b'd');
743            assert_eq!(dt.kind(), b'f');
744            assert_eq!(dt.byteorder(), b'=');
745            assert_eq!(dt.is_native_byteorder(), Some(true));
746            assert_eq!(dt.itemsize(), 8);
747            assert_eq!(dt.alignment(), 8);
748            assert!(!dt.has_object());
749            assert!(dt.names().is_none());
750            assert!(!dt.has_fields());
751            assert!(!dt.is_aligned_struct());
752            assert!(!dt.has_subarray());
753            assert!(dt.base().is_equiv_to(&dt));
754            assert_eq!(dt.ndim(), 0);
755            assert_eq!(dt.shape(), Vec::<usize>::new());
756        });
757    }
758
759    #[test]
760    fn test_dtype_methods_subarray() {
761        Python::attach(|py| {
762            let locals = PyDict::new(py);
763            py_run!(
764                py,
765                *locals,
766                "dtype = __import__('numpy').dtype(('f8', (2, 3)))"
767            );
768            let dt = locals
769                .get_item("dtype")
770                .expect("Operation failed")
771                .expect("Operation failed")
772                .cast_into::<PyArrayDescr>()
773                .expect("Operation failed");
774
775            assert_eq!(dt.num(), NPY_TYPES::NPY_VOID as c_int);
776            assert_eq!(dt.flags(), 0);
777            assert_eq!(dt.typeobj().qualname().expect("Operation failed"), "void");
778            assert_eq!(dt.char(), b'V');
779            assert_eq!(dt.kind(), b'V');
780            assert_eq!(dt.byteorder(), b'|');
781            assert_eq!(dt.is_native_byteorder(), None);
782            assert_eq!(dt.itemsize(), 48);
783            assert_eq!(dt.alignment(), 8);
784            assert!(!dt.has_object());
785            assert!(dt.names().is_none());
786            assert!(!dt.has_fields());
787            assert!(!dt.is_aligned_struct());
788            assert!(dt.has_subarray());
789            assert_eq!(dt.ndim(), 2);
790            assert_eq!(dt.shape(), vec![2, 3]);
791            assert!(dt.base().is_equiv_to(&dtype::<f64>(py)));
792        });
793    }
794
795    #[test]
796    fn test_dtype_methods_record() {
797        Python::attach(|py| {
798            let locals = PyDict::new(py);
799            py_run!(
800                py,
801                *locals,
802                "dtype = __import__('numpy').dtype([('x', 'u1'), ('y', 'f8'), ('z', 'O')], align=True)"
803            );
804            let dt = locals
805                .get_item("dtype")
806                .expect("Operation failed")
807                .expect("Operation failed")
808                .cast_into::<PyArrayDescr>()
809                .expect("Operation failed");
810
811            assert_eq!(dt.num(), NPY_TYPES::NPY_VOID as c_int);
812            assert_ne!(dt.flags() & NPY_ITEM_HASOBJECT, 0);
813            assert_ne!(dt.flags() & NPY_NEEDS_PYAPI, 0);
814            assert_ne!(dt.flags() & NPY_ALIGNED_STRUCT, 0);
815            assert_eq!(dt.typeobj().qualname().expect("Operation failed"), "void");
816            assert_eq!(dt.char(), b'V');
817            assert_eq!(dt.kind(), b'V');
818            assert_eq!(dt.byteorder(), b'|');
819            assert_eq!(dt.is_native_byteorder(), None);
820            assert_eq!(dt.itemsize(), 24);
821            assert_eq!(dt.alignment(), 8);
822            assert!(dt.has_object());
823            assert_eq!(
824                dt.names(),
825                Some(vec!["x".to_owned(), "y".to_owned(), "z".to_owned()])
826            );
827            assert!(dt.has_fields());
828            assert!(dt.is_aligned_struct());
829            assert!(!dt.has_subarray());
830            assert_eq!(dt.ndim(), 0);
831            assert_eq!(dt.shape(), Vec::<usize>::new());
832            assert!(dt.base().is_equiv_to(&dt));
833            let x = dt.get_field("x").expect("Operation failed");
834            assert!(x.0.is_equiv_to(&dtype::<u8>(py)));
835            assert_eq!(x.1, 0);
836            let y = dt.get_field("y").expect("Operation failed");
837            assert!(y.0.is_equiv_to(&dtype::<f64>(py)));
838            assert_eq!(y.1, 8);
839            let z = dt.get_field("z").expect("Operation failed");
840            assert!(z.0.is_equiv_to(&dtype::<Py<PyAny>>(py)));
841            assert_eq!(z.1, 16);
842        });
843    }
844}
scirs2_numpy/dtype.rs

scirs2_numpy/
dtype.rs
