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::{ffi, pyobject_native_type_named, Bound, PyAny, PyTypeInfo, Python};
7
8use crate::array::{PyArray, PyArrayMethods};
9use crate::cold;
10use crate::dtype::PyArrayDescr;
11use crate::npyffi;
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    const NAME: &'static str = "PyUntypedArray";
64    const MODULE: Option<&'static str> = Some("numpy");
65
66    fn type_object_raw<'py>(py: Python<'py>) -> *mut ffi::PyTypeObject {
67        unsafe { npyffi::PY_ARRAY_API.get_type_object(py, npyffi::NpyTypes::PyArray_Type) }
68    }
69
70    fn is_type_of(ob: &Bound<'_, PyAny>) -> bool {
71        unsafe { npyffi::PyArray_Check(ob.py(), ob.as_ptr()) != 0 }
72    }
73}
74
75pyobject_native_type_named!(PyUntypedArray);
76
77/// Implementation of functionality for [`PyUntypedArray`].
78#[doc(alias = "PyUntypedArray")]
79pub trait PyUntypedArrayMethods<'py>: Sealed {
80    /// Returns a raw pointer to the underlying [`PyArrayObject`][npyffi::PyArrayObject].
81    fn as_array_ptr(&self) -> *mut npyffi::PyArrayObject;
82
83    /// Returns the `dtype` of the array.
84    ///
85    /// See also [`ndarray.dtype`][ndarray-dtype] and [`PyArray_DTYPE`][PyArray_DTYPE].
86    ///
87    /// # Example
88    ///
89    /// ```
90    /// use numpy::prelude::*;
91    /// use numpy::{dtype, PyArray};
92    /// use pyo3::Python;
93    ///
94    /// Python::attach(|py| {
95    ///    let array = PyArray::from_vec(py, vec![1_i32, 2, 3]);
96    ///
97    ///    assert!(array.dtype().is_equiv_to(&dtype::<i32>(py)));
98    /// });
99    /// ```
100    ///
101    /// [ndarray-dtype]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.dtype.html
102    /// [PyArray_DTYPE]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_DTYPE
103    fn dtype(&self) -> Bound<'py, PyArrayDescr>;
104
105    /// Returns `true` if the internal data of the array is aligned for the dtype.
106    ///
107    /// Note that NumPy considers zero-length data to be aligned regardless of the base pointer,
108    /// which is a weaker condition than Rust's slice guarantees.  [PyArrayMethods::as_slice] will
109    /// safely handle the case of a misaligned zero-length array.
110    ///
111    /// # Example
112    ///
113    /// ```
114    /// use numpy::{PyArray1, PyUntypedArrayMethods};
115    /// use pyo3::{types::{IntoPyDict, PyAnyMethods}, Python, ffi::c_str};
116    ///
117    /// # fn main() -> pyo3::PyResult<()> {
118    /// Python::attach(|py| {
119    ///     let array = PyArray1::<u16>::zeros(py, 8, false);
120    ///     assert!(array.is_aligned());
121    ///
122    ///     let view = py
123    ///         .eval(
124    ///             c_str!("array.view('u1')[1:-1].view('u2')"),
125    ///             None,
126    ///             Some(&[("array", array)].into_py_dict(py)?),
127    ///         )?
128    ///         .cast_into::<PyArray1<u16>>()?;
129    ///     assert!(!view.is_aligned());
130    /// #   Ok(())
131    /// })
132    /// # }
133    /// ```
134    fn is_aligned(&self) -> bool {
135        unsafe { check_flags(&*self.as_array_ptr(), npyffi::NPY_ARRAY_ALIGNED) }
136    }
137
138    /// Returns `true` if the internal data of the array is contiguous,
139    /// indepedently of whether C-style/row-major or Fortran-style/column-major.
140    ///
141    /// # Example
142    ///
143    /// ```
144    /// use numpy::{PyArray1, PyUntypedArrayMethods};
145    /// use pyo3::{types::{IntoPyDict, PyAnyMethods}, Python, ffi::c_str};
146    ///
147    /// # fn main() -> pyo3::PyResult<()> {
148    /// Python::attach(|py| {
149    ///     let array = PyArray1::arange(py, 0, 10, 1);
150    ///     assert!(array.is_contiguous());
151    ///
152    ///     let view = py
153    ///         .eval(c_str!("array[::2]"), None, Some(&[("array", array)].into_py_dict(py)?))?
154    ///         .cast_into::<PyArray1<i32>>()?;
155    ///     assert!(!view.is_contiguous());
156    /// #   Ok(())
157    /// })
158    /// # }
159    /// ```
160    fn is_contiguous(&self) -> bool {
161        unsafe {
162            check_flags(
163                &*self.as_array_ptr(),
164                npyffi::NPY_ARRAY_C_CONTIGUOUS | npyffi::NPY_ARRAY_F_CONTIGUOUS,
165            )
166        }
167    }
168
169    /// Returns `true` if the internal data of the array is Fortran-style/column-major contiguous.
170    fn is_fortran_contiguous(&self) -> bool {
171        unsafe { check_flags(&*self.as_array_ptr(), npyffi::NPY_ARRAY_F_CONTIGUOUS) }
172    }
173
174    /// Returns `true` if the internal data of the array is C-style/row-major contiguous.
175    fn is_c_contiguous(&self) -> bool {
176        unsafe { check_flags(&*self.as_array_ptr(), npyffi::NPY_ARRAY_C_CONTIGUOUS) }
177    }
178
179    /// Returns the number of dimensions of the array.
180    ///
181    /// See also [`ndarray.ndim`][ndarray-ndim] and [`PyArray_NDIM`][PyArray_NDIM].
182    ///
183    /// # Example
184    ///
185    /// ```
186    /// use numpy::{PyArray3, PyUntypedArrayMethods};
187    /// use pyo3::Python;
188    ///
189    /// Python::attach(|py| {
190    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
191    ///
192    ///     assert_eq!(arr.ndim(), 3);
193    /// });
194    /// ```
195    ///
196    /// [ndarray-ndim]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.ndim.html
197    /// [PyArray_NDIM]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_NDIM
198    #[inline]
199    fn ndim(&self) -> usize {
200        unsafe { (*self.as_array_ptr()).nd as usize }
201    }
202
203    /// Returns a slice indicating how many bytes to advance when iterating along each axis.
204    ///
205    /// See also [`ndarray.strides`][ndarray-strides] and [`PyArray_STRIDES`][PyArray_STRIDES].
206    ///
207    /// # Example
208    ///
209    /// ```
210    /// use numpy::{PyArray3, PyUntypedArrayMethods};
211    /// use pyo3::Python;
212    ///
213    /// Python::attach(|py| {
214    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
215    ///
216    ///     assert_eq!(arr.strides(), &[240, 48, 8]);
217    /// });
218    /// ```
219    /// [ndarray-strides]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.strides.html
220    /// [PyArray_STRIDES]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_STRIDES
221    #[inline]
222    fn strides(&self) -> &[isize] {
223        let n = self.ndim();
224        if n == 0 {
225            cold();
226            return &[];
227        }
228        let ptr = self.as_array_ptr();
229        unsafe {
230            let p = (*ptr).strides;
231            slice::from_raw_parts(p, n)
232        }
233    }
234
235    /// Returns a slice which contains dimmensions of the array.
236    ///
237    /// See also [`ndarray.shape`][ndaray-shape] and [`PyArray_DIMS`][PyArray_DIMS].
238    ///
239    /// # Example
240    ///
241    /// ```
242    /// use numpy::{PyArray3, PyUntypedArrayMethods};
243    /// use pyo3::Python;
244    ///
245    /// Python::attach(|py| {
246    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
247    ///
248    ///     assert_eq!(arr.shape(), &[4, 5, 6]);
249    /// });
250    /// ```
251    ///
252    /// [ndarray-shape]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.shape.html
253    /// [PyArray_DIMS]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_DIMS
254    #[inline]
255    fn shape(&self) -> &[usize] {
256        let n = self.ndim();
257        if n == 0 {
258            cold();
259            return &[];
260        }
261        let ptr = self.as_array_ptr();
262        unsafe {
263            let p = (*ptr).dimensions as *mut usize;
264            slice::from_raw_parts(p, n)
265        }
266    }
267
268    /// Calculates the total number of elements in the array.
269    fn len(&self) -> usize {
270        self.shape().iter().product()
271    }
272
273    /// Returns `true` if the there are no elements in the array.
274    fn is_empty(&self) -> bool {
275        self.shape().contains(&0)
276    }
277}
278
279mod sealed {
280    pub trait Sealed {}
281}
282
283use sealed::Sealed;
284
285fn check_flags(obj: &npyffi::PyArrayObject, flags: i32) -> bool {
286    obj.flags & flags != 0
287}
288
289impl<'py> PyUntypedArrayMethods<'py> for Bound<'py, PyUntypedArray> {
290    #[inline]
291    fn as_array_ptr(&self) -> *mut npyffi::PyArrayObject {
292        self.as_ptr().cast()
293    }
294
295    fn dtype(&self) -> Bound<'py, PyArrayDescr> {
296        unsafe {
297            let descr_ptr = (*self.as_array_ptr()).descr;
298            Bound::from_borrowed_ptr(self.py(), descr_ptr.cast()).cast_into_unchecked()
299        }
300    }
301}
302
303impl Sealed for Bound<'_, PyUntypedArray> {}
304
305// We won't be able to provide a `Deref` impl from `Bound<'_, PyArray<T, D>>` to
306// `Bound<'_, PyUntypedArray>`, so this seems to be the next best thing to do
307impl<'py, T, D> PyUntypedArrayMethods<'py> for Bound<'py, PyArray<T, D>> {
308    #[inline]
309    fn as_array_ptr(&self) -> *mut npyffi::PyArrayObject {
310        self.as_untyped().as_array_ptr()
311    }
312
313    #[inline]
314    fn dtype(&self) -> Bound<'py, PyArrayDescr> {
315        self.as_untyped().dtype()
316    }
317}
318
319impl<T, D> Sealed for Bound<'_, PyArray<T, D>> {}
numpy/untyped_array.rs

numpy/
untyped_array.rs
