numpy/
untyped_array.rs

1//! Safe, untyped interface for NumPy's [N-dimensional arrays][ndarray]
2//!
3//! [ndarray]: https://numpy.org/doc/stable/reference/arrays.ndarray.html
4use std::slice;
5
6use pyo3::{Bound, PyAny, PyTypeInfo, Python, ffi, pyobject_native_type_named};
7
8use crate::array::{PyArray, PyArrayMethods};
9use crate::cold;
10use crate::dtype::PyArrayDescr;
11use crate::npyffi::{self, _PyArray_GET_ITEM_DATA};
12
13/// A safe, untyped wrapper for NumPy's [`ndarray`] class.
14///
15/// Unlike [`PyArray<T,D>`][crate::PyArray], this type does not constrain either element type `T` nor the dimensionality `D`.
16/// This can be useful to inspect function arguments, but it prevents operating on the elements without further downcasts.
17///
18/// When both element type `T` and index type `D` are known, these values can be downcast to `PyArray<T, D>`. In addition,
19/// `PyArray<T, D>` can be dereferenced to a `PyUntypedArray` and can therefore automatically access its methods.
20///
21/// # Example
22///
23/// Taking `PyUntypedArray` can be helpful to implement polymorphic entry points:
24///
25/// ```
26/// # use pyo3::prelude::*;
27/// use pyo3::exceptions::PyTypeError;
28/// use numpy::{Element, PyUntypedArray, PyArray1, dtype};
29/// use numpy::{PyUntypedArrayMethods, PyArrayMethods, PyArrayDescrMethods};
30///
31/// #[pyfunction]
32/// fn entry_point(py: Python<'_>, array: &Bound<'_, PyUntypedArray>) -> PyResult<()> {
33///     fn implementation<T: Element>(array: &Bound<'_, PyArray1<T>>) -> PyResult<()> {
34///         /* .. */
35///
36///         Ok(())
37///     }
38///
39///     let element_type = array.dtype();
40///
41///     if element_type.is_equiv_to(&dtype::<f32>(py)) {
42///         let array = array.cast::<PyArray1<f32>>()?;
43///
44///         implementation(array)
45///     } else if element_type.is_equiv_to(&dtype::<f64>(py)) {
46///         let array = array.cast::<PyArray1<f64>>()?;
47///
48///         implementation(array)
49///     } else {
50///         Err(PyTypeError::new_err(format!("Unsupported element type: {}", element_type)))
51///     }
52/// }
53/// #
54/// # Python::attach(|py| {
55/// #   let array = PyArray1::<f64>::zeros(py, 42, false);
56/// #   entry_point(py, array.as_untyped())
57/// # }).unwrap();
58/// ```
59#[repr(transparent)]
60pub struct PyUntypedArray(PyAny);
61
62unsafe impl PyTypeInfo for PyUntypedArray {
63    fn type_object_raw<'py>(py: Python<'py>) -> *mut ffi::PyTypeObject {
64        unsafe { npyffi::get_type_object(py, npyffi::NpyTypes::PyArray_Type) }
65    }
66
67    fn is_type_of(ob: &Bound<'_, PyAny>) -> bool {
68        unsafe { npyffi::PyArray_Check(ob.py(), ob.as_ptr()) != 0 }
69    }
70}
71
72pyobject_native_type_named!(PyUntypedArray);
73
74/// Implementation of functionality for [`PyUntypedArray`].
75#[doc(alias = "PyUntypedArray")]
76pub trait PyUntypedArrayMethods<'py>: Sealed {
77    /// Returns a raw pointer to the underlying [`PyArrayObject`][npyffi::PyArrayObject].
78    fn as_array_ptr(&self) -> *mut npyffi::PyArrayObject;
79
80    /// Returns the `dtype` of the array.
81    ///
82    /// See also [`ndarray.dtype`][ndarray-dtype] and [`PyArray_DTYPE`][PyArray_DTYPE].
83    ///
84    /// # Example
85    ///
86    /// ```
87    /// use numpy::prelude::*;
88    /// use numpy::{dtype, PyArray};
89    /// use pyo3::Python;
90    ///
91    /// Python::attach(|py| {
92    ///    let array = PyArray::from_vec(py, vec![1_i32, 2, 3]);
93    ///
94    ///    assert!(array.dtype().is_equiv_to(&dtype::<i32>(py)));
95    /// });
96    /// ```
97    ///
98    /// [ndarray-dtype]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.dtype.html
99    /// [PyArray_DTYPE]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_DTYPE
100    fn dtype(&self) -> Bound<'py, PyArrayDescr>;
101
102    /// Returns `true` if the internal data of the array is aligned for the dtype.
103    ///
104    /// Note that NumPy considers zero-length data to be aligned regardless of the base pointer,
105    /// which is a weaker condition than Rust's slice guarantees.  [PyArrayMethods::as_slice] will
106    /// safely handle the case of a misaligned zero-length array.
107    ///
108    /// # Example
109    ///
110    /// ```
111    /// use numpy::{PyArray1, PyUntypedArrayMethods};
112    /// use pyo3::{types::{IntoPyDict, PyAnyMethods}, Python};
113    ///
114    /// # fn main() -> pyo3::PyResult<()> {
115    /// Python::attach(|py| {
116    ///     let array = PyArray1::<u16>::zeros(py, 8, false);
117    ///     assert!(array.is_aligned());
118    ///
119    ///     let view = py
120    ///         .eval(
121    ///             c"array.view('u1')[1:-1].view('u2')",
122    ///             None,
123    ///             Some(&[("array", array)].into_py_dict(py)?),
124    ///         )?
125    ///         .cast_into::<PyArray1<u16>>()?;
126    ///     assert!(!view.is_aligned());
127    /// #   Ok(())
128    /// })
129    /// # }
130    /// ```
131    fn is_aligned(&self) -> bool {
132        unsafe { check_flags(&*self.as_array_ptr(), npyffi::NPY_ARRAY_ALIGNED) }
133    }
134
135    /// Returns `true` if the internal data of the array is contiguous,
136    /// independently of whether C-style/row-major or Fortran-style/column-major.
137    ///
138    /// # Example
139    ///
140    /// ```
141    /// use numpy::{PyArray1, PyUntypedArrayMethods};
142    /// use pyo3::{types::{IntoPyDict, PyAnyMethods}, Python};
143    ///
144    /// # fn main() -> pyo3::PyResult<()> {
145    /// Python::attach(|py| {
146    ///     let array = PyArray1::arange(py, 0, 10, 1);
147    ///     assert!(array.is_contiguous());
148    ///
149    ///     let view = py
150    ///         .eval(c"array[::2]", None, Some(&[("array", array)].into_py_dict(py)?))?
151    ///         .cast_into::<PyArray1<i32>>()?;
152    ///     assert!(!view.is_contiguous());
153    /// #   Ok(())
154    /// })
155    /// # }
156    /// ```
157    fn is_contiguous(&self) -> bool {
158        unsafe {
159            check_flags(
160                &*self.as_array_ptr(),
161                npyffi::NPY_ARRAY_C_CONTIGUOUS | npyffi::NPY_ARRAY_F_CONTIGUOUS,
162            )
163        }
164    }
165
166    /// Returns `true` if the internal data of the array is Fortran-style/column-major contiguous.
167    fn is_fortran_contiguous(&self) -> bool {
168        unsafe { check_flags(&*self.as_array_ptr(), npyffi::NPY_ARRAY_F_CONTIGUOUS) }
169    }
170
171    /// Returns `true` if the internal data of the array is C-style/row-major contiguous.
172    fn is_c_contiguous(&self) -> bool {
173        unsafe { check_flags(&*self.as_array_ptr(), npyffi::NPY_ARRAY_C_CONTIGUOUS) }
174    }
175
176    /// Returns the number of dimensions of the array.
177    ///
178    /// See also [`ndarray.ndim`][ndarray-ndim] and [`PyArray_NDIM`][PyArray_NDIM].
179    ///
180    /// # Example
181    ///
182    /// ```
183    /// use numpy::{PyArray3, PyUntypedArrayMethods};
184    /// use pyo3::Python;
185    ///
186    /// Python::attach(|py| {
187    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
188    ///
189    ///     assert_eq!(arr.ndim(), 3);
190    /// });
191    /// ```
192    ///
193    /// [ndarray-ndim]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.ndim.html
194    /// [PyArray_NDIM]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_NDIM
195    #[inline]
196    fn ndim(&self) -> usize {
197        unsafe { (*_PyArray_GET_ITEM_DATA(self.as_array_ptr().cast_const())).nd as usize }
198    }
199
200    /// Returns a slice indicating how many bytes to advance when iterating along each axis.
201    ///
202    /// See also [`ndarray.strides`][ndarray-strides] and [`PyArray_STRIDES`][PyArray_STRIDES].
203    ///
204    /// # Example
205    ///
206    /// ```
207    /// use numpy::{PyArray3, PyUntypedArrayMethods};
208    /// use pyo3::Python;
209    ///
210    /// Python::attach(|py| {
211    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
212    ///
213    ///     assert_eq!(arr.strides(), &[240, 48, 8]);
214    /// });
215    /// ```
216    /// [ndarray-strides]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.strides.html
217    /// [PyArray_STRIDES]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_STRIDES
218    #[inline]
219    fn strides(&self) -> &[isize] {
220        let n = self.ndim();
221        if n == 0 {
222            cold();
223            return &[];
224        }
225        let ptr = unsafe { _PyArray_GET_ITEM_DATA(self.as_array_ptr().cast_const()) };
226        unsafe {
227            let p = (*ptr).strides;
228            slice::from_raw_parts(p, n)
229        }
230    }
231
232    /// Returns a slice which contains dimensions of the array.
233    ///
234    /// See also [`ndarray.shape`][ndarray-shape] and [`PyArray_DIMS`][PyArray_DIMS].
235    ///
236    /// # Example
237    ///
238    /// ```
239    /// use numpy::{PyArray3, PyUntypedArrayMethods};
240    /// use pyo3::Python;
241    ///
242    /// Python::attach(|py| {
243    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
244    ///
245    ///     assert_eq!(arr.shape(), &[4, 5, 6]);
246    /// });
247    /// ```
248    ///
249    /// [ndarray-shape]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.shape.html
250    /// [PyArray_DIMS]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_DIMS
251    #[inline]
252    fn shape(&self) -> &[usize] {
253        let n = self.ndim();
254        if n == 0 {
255            cold();
256            return &[];
257        }
258        let ptr = unsafe { _PyArray_GET_ITEM_DATA(self.as_array_ptr().cast_const()) };
259        unsafe {
260            let p = (*ptr).dimensions as *mut usize;
261            slice::from_raw_parts(p, n)
262        }
263    }
264
265    /// Calculates the total number of elements in the array.
266    fn len(&self) -> usize {
267        self.shape().iter().product()
268    }
269
270    /// Returns `true` if there are no elements in the array.
271    fn is_empty(&self) -> bool {
272        self.shape().contains(&0)
273    }
274}
275
276mod sealed {
277    pub trait Sealed {}
278}
279
280use sealed::Sealed;
281
282fn check_flags(obj: &npyffi::PyArrayObject, flags: i32) -> bool {
283    unsafe { (*_PyArray_GET_ITEM_DATA(obj)).flags & flags != 0 }
284}
285
286impl<'py> PyUntypedArrayMethods<'py> for Bound<'py, PyUntypedArray> {
287    #[inline]
288    fn as_array_ptr(&self) -> *mut npyffi::PyArrayObject {
289        self.as_ptr().cast()
290    }
291
292    fn dtype(&self) -> Bound<'py, PyArrayDescr> {
293        unsafe {
294            let descr_ptr =
295                (*npyffi::_PyArray_GET_ITEM_DATA(self.as_array_ptr().cast_const())).descr;
296            Bound::from_borrowed_ptr(self.py(), descr_ptr.cast()).cast_into_unchecked()
297        }
298    }
299}
300
301impl Sealed for Bound<'_, PyUntypedArray> {}
302
303// We won't be able to provide a `Deref` impl from `Bound<'_, PyArray<T, D>>` to
304// `Bound<'_, PyUntypedArray>`, so this seems to be the next best thing to do
305impl<'py, T, D> PyUntypedArrayMethods<'py> for Bound<'py, PyArray<T, D>> {
306    #[inline]
307    fn as_array_ptr(&self) -> *mut npyffi::PyArrayObject {
308        self.as_untyped().as_array_ptr()
309    }
310
311    #[inline]
312    fn dtype(&self) -> Bound<'py, PyArrayDescr> {
313        self.as_untyped().dtype()
314    }
315}
316
317impl<T, D> Sealed for Bound<'_, PyArray<T, D>> {}
numpy/untyped_array.rs

numpy/
untyped_array.rs
