numpy/
dtype.rs

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

numpy/
dtype.rs
