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 contiguous,
106    /// indepedently of whether C-style/row-major or Fortran-style/column-major.
107    ///
108    /// # Example
109    ///
110    /// ```
111    /// use numpy::{PyArray1, PyUntypedArrayMethods};
112    /// use pyo3::{types::{IntoPyDict, PyAnyMethods}, Python, ffi::c_str};
113    ///
114    /// # fn main() -> pyo3::PyResult<()> {
115    /// Python::attach(|py| {
116    ///     let array = PyArray1::arange(py, 0, 10, 1);
117    ///     assert!(array.is_contiguous());
118    ///
119    ///     let view = py
120    ///         .eval(c_str!("array[::2]"), None, Some(&[("array", array)].into_py_dict(py)?))?
121    ///         .cast_into::<PyArray1<i32>>()?;
122    ///     assert!(!view.is_contiguous());
123    /// #   Ok(())
124    /// })
125    /// # }
126    /// ```
127    fn is_contiguous(&self) -> bool {
128        unsafe {
129            check_flags(
130                &*self.as_array_ptr(),
131                npyffi::NPY_ARRAY_C_CONTIGUOUS | npyffi::NPY_ARRAY_F_CONTIGUOUS,
132            )
133        }
134    }
135
136    /// Returns `true` if the internal data of the array is Fortran-style/column-major contiguous.
137    fn is_fortran_contiguous(&self) -> bool {
138        unsafe { check_flags(&*self.as_array_ptr(), npyffi::NPY_ARRAY_F_CONTIGUOUS) }
139    }
140
141    /// Returns `true` if the internal data of the array is C-style/row-major contiguous.
142    fn is_c_contiguous(&self) -> bool {
143        unsafe { check_flags(&*self.as_array_ptr(), npyffi::NPY_ARRAY_C_CONTIGUOUS) }
144    }
145
146    /// Returns the number of dimensions of the array.
147    ///
148    /// See also [`ndarray.ndim`][ndarray-ndim] and [`PyArray_NDIM`][PyArray_NDIM].
149    ///
150    /// # Example
151    ///
152    /// ```
153    /// use numpy::{PyArray3, PyUntypedArrayMethods};
154    /// use pyo3::Python;
155    ///
156    /// Python::attach(|py| {
157    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
158    ///
159    ///     assert_eq!(arr.ndim(), 3);
160    /// });
161    /// ```
162    ///
163    /// [ndarray-ndim]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.ndim.html
164    /// [PyArray_NDIM]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_NDIM
165    #[inline]
166    fn ndim(&self) -> usize {
167        unsafe { (*self.as_array_ptr()).nd as usize }
168    }
169
170    /// Returns a slice indicating how many bytes to advance when iterating along each axis.
171    ///
172    /// See also [`ndarray.strides`][ndarray-strides] and [`PyArray_STRIDES`][PyArray_STRIDES].
173    ///
174    /// # Example
175    ///
176    /// ```
177    /// use numpy::{PyArray3, PyUntypedArrayMethods};
178    /// use pyo3::Python;
179    ///
180    /// Python::attach(|py| {
181    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
182    ///
183    ///     assert_eq!(arr.strides(), &[240, 48, 8]);
184    /// });
185    /// ```
186    /// [ndarray-strides]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.strides.html
187    /// [PyArray_STRIDES]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_STRIDES
188    #[inline]
189    fn strides(&self) -> &[isize] {
190        let n = self.ndim();
191        if n == 0 {
192            cold();
193            return &[];
194        }
195        let ptr = self.as_array_ptr();
196        unsafe {
197            let p = (*ptr).strides;
198            slice::from_raw_parts(p, n)
199        }
200    }
201
202    /// Returns a slice which contains dimmensions of the array.
203    ///
204    /// See also [`ndarray.shape`][ndaray-shape] and [`PyArray_DIMS`][PyArray_DIMS].
205    ///
206    /// # Example
207    ///
208    /// ```
209    /// use numpy::{PyArray3, PyUntypedArrayMethods};
210    /// use pyo3::Python;
211    ///
212    /// Python::attach(|py| {
213    ///     let arr = PyArray3::<f64>::zeros(py, [4, 5, 6], false);
214    ///
215    ///     assert_eq!(arr.shape(), &[4, 5, 6]);
216    /// });
217    /// ```
218    ///
219    /// [ndarray-shape]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.shape.html
220    /// [PyArray_DIMS]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_DIMS
221    #[inline]
222    fn shape(&self) -> &[usize] {
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).dimensions as *mut usize;
231            slice::from_raw_parts(p, n)
232        }
233    }
234
235    /// Calculates the total number of elements in the array.
236    fn len(&self) -> usize {
237        self.shape().iter().product()
238    }
239
240    /// Returns `true` if the there are no elements in the array.
241    fn is_empty(&self) -> bool {
242        self.shape().contains(&0)
243    }
244}
245
246mod sealed {
247    pub trait Sealed {}
248}
249
250use sealed::Sealed;
251
252fn check_flags(obj: &npyffi::PyArrayObject, flags: i32) -> bool {
253    obj.flags & flags != 0
254}
255
256impl<'py> PyUntypedArrayMethods<'py> for Bound<'py, PyUntypedArray> {
257    #[inline]
258    fn as_array_ptr(&self) -> *mut npyffi::PyArrayObject {
259        self.as_ptr().cast()
260    }
261
262    fn dtype(&self) -> Bound<'py, PyArrayDescr> {
263        unsafe {
264            let descr_ptr = (*self.as_array_ptr()).descr;
265            Bound::from_borrowed_ptr(self.py(), descr_ptr.cast()).cast_into_unchecked()
266        }
267    }
268}
269
270impl Sealed for Bound<'_, PyUntypedArray> {}
271
272// We won't be able to provide a `Deref` impl from `Bound<'_, PyArray<T, D>>` to
273// `Bound<'_, PyUntypedArray>`, so this seems to be the next best thing to do
274impl<'py, T, D> PyUntypedArrayMethods<'py> for Bound<'py, PyArray<T, D>> {
275    #[inline]
276    fn as_array_ptr(&self) -> *mut npyffi::PyArrayObject {
277        self.as_untyped().as_array_ptr()
278    }
279
280    #[inline]
281    fn dtype(&self) -> Bound<'py, PyArrayDescr> {
282        self.as_untyped().dtype()
283    }
284}
285
286impl<T, D> Sealed for Bound<'_, PyArray<T, D>> {}
numpy/untyped_array.rs

numpy/
untyped_array.rs
