pyo3_arraylike/
lib.rs

1//! This crate provides a single struct `PyArrayLike<T,D>` which can be used for extracting an array from any Python object which can be regarded as an array of type `T` and dimension `D` in a reasonable way.
2
3#![deny(missing_docs, missing_debug_implementations)]
4
5#[cfg(test)]
6mod test;
7
8use ndarray::{Array, ArrayView, Axis, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
9use numpy::{
10    ndarray::Dimension,
11    pyo3::{
12        exceptions::PyValueError, types::PyAnyMethods, Borrowed, Bound, FromPyObject, PyAny, PyErr,
13        PyResult, Python,
14    },
15    Element, IntoPyArray, PyArray, PyArrayMethods, PyReadonlyArray,
16};
17use std::fmt::Debug;
18
19/// To be used for extracting an array from any Python object which can be regarded as an array of type `T` and dimension `D` in a reasonable way.
20#[derive(Debug)]
21pub struct PyArrayLike<'py, T, D>(ArrayLike<'py, T, D>)
22where
23    T: Element,
24    D: Dimension;
25
26enum ArrayLike<'py, T, D>
27where
28    T: Element,
29    D: Dimension,
30{
31    PyRef(PyReadonlyArray<'py, T, D>),
32    Owned(Array<T, D>, Python<'py>),
33}
34
35impl<'py, T, D> Debug for ArrayLike<'py, T, D>
36where
37    T: Element + Debug,
38    D: Dimension,
39{
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        match self {
42            Self::PyRef(py_array) => f.debug_tuple("PyRef").field(py_array).finish(),
43            Self::Owned(array, _) => f.debug_tuple("Owned").field(array).finish(),
44        }
45    }
46}
47
48impl<'py, T, D> PyArrayLike<'py, T, D>
49where
50    T: Element,
51    D: Dimension,
52{
53    /// Consumes `self` and moves its data into an owned array.
54    pub fn into_owned_array(self) -> Array<T, D> {
55        match self.0 {
56            ArrayLike::PyRef(py_array) => py_array.to_owned_array(),
57            ArrayLike::Owned(array, _) => array,
58        }
59    }
60
61    /// Consumes `self` and moves its data into a numpy array.
62    pub fn into_pyarray(self) -> PyReadonlyArray<'py, T, D> {
63        match self.0 {
64            ArrayLike::PyRef(py_array) => py_array,
65            ArrayLike::Owned(array, py) => array.into_pyarray(py).readonly(),
66        }
67    }
68
69    /// Return a read-only view of the array.
70    pub fn view<'s>(&'s self) -> ArrayView<'s, T, D> {
71        match &self.0 {
72            ArrayLike::PyRef(py_array) => py_array.as_array(),
73            ArrayLike::Owned(array, _) => array.view(),
74        }
75    }
76
77    /// Return the array’s data as a slice, if it is contiguous and in standard order.
78    pub fn as_slice(&self) -> Option<&[T]> {
79        match &self.0 {
80            ArrayLike::PyRef(py_array) => py_array.as_slice().ok(),
81            ArrayLike::Owned(array, _) => array.as_slice(),
82        }
83    }
84
85    /// Return the array’s dimension
86    pub fn dim(&self) -> D::Pattern {
87        match &self.0 {
88            ArrayLike::PyRef(py_array) => py_array.dims().into_pattern(),
89            ArrayLike::Owned(array, _) => array.dim(),
90        }
91    }
92}
93
94impl<'py, T, D> From<PyArrayLike<'py, T, D>> for PyReadonlyArray<'py, T, D>
95where
96    T: Element,
97    D: Dimension,
98{
99    fn from(value: PyArrayLike<'py, T, D>) -> Self {
100        value.into_pyarray()
101    }
102}
103
104impl<T, D> From<PyArrayLike<'_, T, D>> for Array<T, D>
105where
106    T: Element,
107    D: Dimension,
108{
109    fn from(value: PyArrayLike<T, D>) -> Self {
110        value.into_owned_array()
111    }
112}
113
114impl<'py, T, D> PyArrayLike<'py, T, D>
115where
116    T: Clone + Element + 'static + for<'a> FromPyObject<'a, 'py>,
117    D: Dimension + 'static,
118{
119    fn from_python(ob: &Bound<'py, PyAny>) -> Option<Self> {
120        if let Ok(array) = ob.cast::<PyArray<T, D>>() {
121            return Some(PyArrayLike(ArrayLike::PyRef(array.readonly())));
122        }
123
124        if matches!(D::NDIM, None | Some(0)) {
125            if let Ok(value) = ob.extract::<T>() {
126                let res = Array::from_elem((), value).into_dimensionality().ok()?;
127                return Some(PyArrayLike(ArrayLike::Owned(res, ob.py())));
128            }
129        }
130
131        if matches!(D::NDIM, None | Some(1)) {
132            if let Ok(array) = ob.extract::<Vec<T>>() {
133                let res = Array::from_vec(array).into_dimensionality().ok()?;
134                return Some(PyArrayLike(ArrayLike::Owned(res, ob.py())));
135            }
136        }
137
138        let sub_arrays = ob
139            .try_iter()
140            .ok()?
141            .map(|item| {
142                item.ok()
143                    .and_then(|ob| <PyArrayLike<T, D::Smaller>>::from_python(&ob))
144            })
145            .collect::<Option<Vec<_>>>()?;
146        let sub_array_views = sub_arrays.iter().map(|x| x.view()).collect::<Vec<_>>();
147        let array = ndarray::stack(Axis(0), &sub_array_views)
148            .ok()?
149            .into_dimensionality()
150            .ok()?;
151        Some(PyArrayLike(ArrayLike::Owned(array, ob.py())))
152    }
153}
154
155impl<'py, T, D> FromPyObject<'_, 'py> for PyArrayLike<'py, T, D>
156where
157    T: Clone + Element + 'static + for<'a> FromPyObject<'a, 'py>,
158    D: Dimension + 'static,
159{
160    type Error = PyErr;
161
162    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {
163        Self::from_python(&ob).ok_or_else(|| {
164            let dtype = T::get_dtype(ob.py());
165            let err_text = match D::NDIM {
166                Some(dim) => format!("Expected an array like of dimension {} containing elements which can be safely casted to {}.", dim, dtype),
167                None => format!("Expected an array like of arbitrary dimension containing elements which can be safely casted to {}.", dtype)
168            };
169            PyValueError::new_err(err_text)})
170    }
171}
172
173/// Zero-dimensional array like.
174pub type PyArrayLike0<'py, T> = PyArrayLike<'py, T, Ix0>;
175/// One-dimensional array like.
176pub type PyArrayLike1<'py, T> = PyArrayLike<'py, T, Ix1>;
177/// Two-dimensional array like.
178pub type PyArrayLike2<'py, T> = PyArrayLike<'py, T, Ix2>;
179/// Three-dimensional array like.
180pub type PyArrayLike3<'py, T> = PyArrayLike<'py, T, Ix3>;
181/// Four-dimensional array like.
182pub type PyArrayLike4<'py, T> = PyArrayLike<'py, T, Ix4>;
183/// Five-dimensional array like.
184pub type PyArrayLike5<'py, T> = PyArrayLike<'py, T, Ix5>;
185/// Six-dimensional array like.
186pub type PyArrayLike6<'py, T> = PyArrayLike<'py, T, Ix6>;
187/// Array like of any dimension.
188pub type PyArrayLikeDyn<'py, T> = PyArrayLike<'py, T, IxDyn>;
pyo3_arraylike/lib.rs

pyo3_arraylike/
lib.rs
