numpy/
array_like.rs

1use std::marker::PhantomData;
2use std::ops::Deref;
3
4use ndarray::{Array1, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
5use pyo3::{types::PyAnyMethods, Borrowed, FromPyObject, PyAny, PyErr, PyResult};
6
7use crate::npyffi::NPY_ARRAY_FORCECAST;
8use crate::{array::PyArrayMethods, PY_ARRAY_API};
9use crate::{Element, IntoPyArray, PyArray, PyReadonlyArray, PyUntypedArray};
10
11pub trait Coerce: Sealed {
12    const ALLOW_TYPE_CHANGE: bool;
13}
14
15mod sealed {
16    pub trait Sealed {}
17}
18
19use sealed::Sealed;
20
21/// Marker type to indicate that the element type received via [`PyArrayLike`] must match the specified type exactly.
22#[derive(Debug)]
23pub struct TypeMustMatch;
24
25impl Sealed for TypeMustMatch {}
26
27impl Coerce for TypeMustMatch {
28    const ALLOW_TYPE_CHANGE: bool = false;
29}
30
31/// Marker type to indicate that the element type received via [`PyArrayLike`] can be cast to the specified type by NumPy's [`asarray`](https://numpy.org/doc/stable/reference/generated/numpy.asarray.html).
32#[derive(Debug)]
33pub struct AllowTypeChange;
34
35impl Sealed for AllowTypeChange {}
36
37impl Coerce for AllowTypeChange {
38    const ALLOW_TYPE_CHANGE: bool = true;
39}
40
41/// Receiver for arrays or array-like types.
42///
43/// When building API using NumPy in Python, it is common for functions to additionally accept any array-like type such as `list[float]` as arguments.
44/// `PyArrayLike` enables the same pattern in Rust extensions, i.e. by taking this type as the argument of a `#[pyfunction]`,
45/// one will always get access to a [`PyReadonlyArray`] that will either reference to the NumPy array originally passed into the function
46/// or a temporary one created by converting the input type into a NumPy array.
47///
48/// Depending on whether [`TypeMustMatch`] or [`AllowTypeChange`] is used for the `C` type parameter,
49/// the element type must either match the specific type `T` exactly or will be cast to it by NumPy's [`asarray`](https://numpy.org/doc/stable/reference/generated/numpy.asarray.html).
50///
51/// # Example
52///
53/// `PyArrayLike1<'py, T, TypeMustMatch>` will enable you to receive both NumPy arrays and sequences
54///
55/// ```rust
56/// # use pyo3::prelude::*;
57/// use pyo3::py_run;
58/// use numpy::{get_array_module, PyArrayLike1, TypeMustMatch};
59///
60/// #[pyfunction]
61/// fn sum_up<'py>(py: Python<'py>, array: PyArrayLike1<'py, f64, TypeMustMatch>) -> f64 {
62///     array.as_array().sum()
63/// }
64///
65/// Python::attach(|py| {
66///     let np = get_array_module(py).unwrap();
67///     let sum_up = wrap_pyfunction!(sum_up)(py).unwrap();
68///
69///     py_run!(py, np sum_up, r"assert sum_up(np.array([1., 2., 3.])) == 6.");
70///     py_run!(py, np sum_up, r"assert sum_up((1., 2., 3.)) == 6.");
71/// });
72/// ```
73///
74/// but it will not cast the element type if that is required
75///
76/// ```rust,should_panic
77/// use pyo3::prelude::*;
78/// use pyo3::py_run;
79/// use numpy::{get_array_module, PyArrayLike1, TypeMustMatch};
80///
81/// #[pyfunction]
82/// fn sum_up<'py>(py: Python<'py>, array: PyArrayLike1<'py, i32, TypeMustMatch>) -> i32 {
83///     array.as_array().sum()
84/// }
85///
86/// Python::attach(|py| {
87///     let np = get_array_module(py).unwrap();
88///     let sum_up = wrap_pyfunction!(sum_up)(py).unwrap();
89///
90///     py_run!(py, np sum_up, r"assert sum_up(np.array([1., 2., 3.])) == 6");
91/// });
92/// ```
93///
94/// whereas `PyArrayLike1<'py, T, AllowTypeChange>` will do even at the cost loosing precision
95///
96/// ```rust
97/// use pyo3::prelude::*;
98/// use pyo3::py_run;
99/// use numpy::{get_array_module, AllowTypeChange, PyArrayLike1};
100///
101/// #[pyfunction]
102/// fn sum_up<'py>(py: Python<'py>, array: PyArrayLike1<'py, i32, AllowTypeChange>) -> i32 {
103///     array.as_array().sum()
104/// }
105///
106/// Python::attach(|py| {
107///     let np = get_array_module(py).unwrap();
108///     let sum_up = wrap_pyfunction!(sum_up)(py).unwrap();
109///
110///     py_run!(py, np sum_up, r"assert sum_up(np.array([1.5, 2.5])) == 3");
111///     py_run!(py, np sum_up, r"assert sum_up((1.5, 2.5)) == 3");
112/// });
113/// ```
114#[derive(Debug)]
115#[repr(transparent)]
116pub struct PyArrayLike<'py, T, D, C = TypeMustMatch>(PyReadonlyArray<'py, T, D>, PhantomData<C>)
117where
118    T: Element,
119    D: Dimension,
120    C: Coerce;
121
122impl<'py, T, D, C> Deref for PyArrayLike<'py, T, D, C>
123where
124    T: Element,
125    D: Dimension,
126    C: Coerce,
127{
128    type Target = PyReadonlyArray<'py, T, D>;
129
130    fn deref(&self) -> &Self::Target {
131        &self.0
132    }
133}
134
135impl<'a, 'py, T, D, C> FromPyObject<'a, 'py> for PyArrayLike<'py, T, D, C>
136where
137    T: Element + 'py,
138    D: Dimension + 'py,
139    C: Coerce,
140    Vec<T>: FromPyObject<'a, 'py>,
141{
142    type Error = PyErr;
143
144    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
145        if let Ok(array) = ob.cast::<PyArray<T, D>>() {
146            return Ok(Self(array.readonly(), PhantomData));
147        }
148
149        let py = ob.py();
150
151        // If the input is already an ndarray and `TypeMustMatch` is used then no type conversion
152        // should be performed.
153        if (C::ALLOW_TYPE_CHANGE || ob.cast::<PyUntypedArray>().is_err())
154            && matches!(D::NDIM, Some(1))
155        {
156            if let Ok(vec) = ob.extract::<Vec<T>>() {
157                let array = Array1::from(vec)
158                    .into_dimensionality()
159                    .expect("D being compatible to Ix1")
160                    .into_pyarray(py)
161                    .readonly();
162                return Ok(Self(array, PhantomData));
163            }
164        }
165
166        let (dtype, flags) = if C::ALLOW_TYPE_CHANGE || ob.cast::<PyUntypedArray>().is_err() {
167            (Some(T::get_dtype(py)), NPY_ARRAY_FORCECAST)
168        } else {
169            (None, 0)
170        };
171
172        let newtype = dtype
173            .map(|dt| dt.into_ptr().cast())
174            .unwrap_or_else(std::ptr::null_mut);
175
176        let array = unsafe {
177            let ptr = PY_ARRAY_API.PyArray_FromAny(
178                py,
179                ob.as_ptr(),
180                newtype,
181                0,
182                0,
183                flags,
184                std::ptr::null_mut(),
185            );
186
187            pyo3::Bound::from_owned_ptr_or_err(py, ptr)?
188        };
189
190        Ok(Self(array.extract()?, PhantomData))
191    }
192}
193
194/// Receiver for zero-dimensional arrays or array-like types.
195pub type PyArrayLike0<'py, T, C = TypeMustMatch> = PyArrayLike<'py, T, Ix0, C>;
196
197/// Receiver for one-dimensional arrays or array-like types.
198pub type PyArrayLike1<'py, T, C = TypeMustMatch> = PyArrayLike<'py, T, Ix1, C>;
199
200/// Receiver for two-dimensional arrays or array-like types.
201pub type PyArrayLike2<'py, T, C = TypeMustMatch> = PyArrayLike<'py, T, Ix2, C>;
202
203/// Receiver for three-dimensional arrays or array-like types.
204pub type PyArrayLike3<'py, T, C = TypeMustMatch> = PyArrayLike<'py, T, Ix3, C>;
205
206/// Receiver for four-dimensional arrays or array-like types.
207pub type PyArrayLike4<'py, T, C = TypeMustMatch> = PyArrayLike<'py, T, Ix4, C>;
208
209/// Receiver for five-dimensional arrays or array-like types.
210pub type PyArrayLike5<'py, T, C = TypeMustMatch> = PyArrayLike<'py, T, Ix5, C>;
211
212/// Receiver for six-dimensional arrays or array-like types.
213pub type PyArrayLike6<'py, T, C = TypeMustMatch> = PyArrayLike<'py, T, Ix6, C>;
214
215/// Receiver for arrays or array-like types whose dimensionality is determined at runtime.
216pub type PyArrayLikeDyn<'py, T, C = TypeMustMatch> = PyArrayLike<'py, T, IxDyn, C>;
numpy/array_like.rs

numpy/
array_like.rs
