numpy/
convert.rs

1//! Defines conversion traits between Rust types and NumPy data types.
2
3use std::{mem, os::raw::c_int, ptr};
4
5use ndarray::{ArrayBase, Data, Dim, Dimension, IntoDimension, Ix1, OwnedRepr};
6use pyo3::{Bound, Python};
7
8use crate::array::{PyArray, PyArrayMethods};
9use crate::dtype::Element;
10use crate::error::MAX_DIMENSIONALITY_ERR;
11use crate::npyffi::{self, npy_intp};
12use crate::slice_container::PySliceContainer;
13
14/// Conversion trait from owning Rust types into [`PyArray`].
15///
16/// This trait takes ownership of `self`, which means it holds a pointer into the Rust heap.
17///
18/// In addition, some destructive methods like `resize` cannot be used with NumPy arrays constructed using this trait.
19///
20/// # Example
21///
22/// ```
23/// use numpy::{PyArray, IntoPyArray, PyArrayMethods};
24/// use pyo3::Python;
25///
26/// Python::with_gil(|py| {
27///     let py_array = vec![1, 2, 3].into_pyarray(py);
28///
29///     assert_eq!(py_array.readonly().as_slice().unwrap(), &[1, 2, 3]);
30///
31///     // Array cannot be resized when its data is owned by Rust.
32///     unsafe {
33///         assert!(py_array.resize(100).is_err());
34///     }
35/// });
36/// ```
37pub trait IntoPyArray: Sized {
38    /// The element type of resulting array.
39    type Item: Element;
40    /// The dimension type of the resulting array.
41    type Dim: Dimension;
42
43    /// Consumes `self` and moves its data into a NumPy array.
44    fn into_pyarray<'py>(self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>>;
45
46    /// Deprecated name for [`IntoPyArray::into_pyarray`].
47    #[deprecated(since = "0.23.0", note = "renamed to `IntoPyArray::into_pyarray`")]
48    #[inline]
49    fn into_pyarray_bound<'py>(
50        self,
51        py: Python<'py>,
52    ) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
53        self.into_pyarray(py)
54    }
55}
56
57impl<T: Element> IntoPyArray for Box<[T]> {
58    type Item = T;
59    type Dim = Ix1;
60
61    fn into_pyarray<'py>(self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
62        let container = PySliceContainer::from(self);
63        let dims = Dim([container.len]);
64        let strides = [mem::size_of::<T>() as npy_intp];
65        // The data pointer is derived only after dissolving `Box` into `PySliceContainer`
66        // to avoid unsound aliasing of Box<[T]> which is currently noalias,
67        // c.f. https://github.com/rust-lang/unsafe-code-guidelines/issues/326
68        let data_ptr = container.ptr as *mut T;
69        unsafe { PyArray::from_raw_parts(py, dims, strides.as_ptr(), data_ptr, container) }
70    }
71}
72
73impl<T: Element> IntoPyArray for Vec<T> {
74    type Item = T;
75    type Dim = Ix1;
76
77    fn into_pyarray<'py>(mut self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
78        let dims = Dim([self.len()]);
79        let strides = [mem::size_of::<T>() as npy_intp];
80        let data_ptr = self.as_mut_ptr();
81        unsafe {
82            PyArray::from_raw_parts(
83                py,
84                dims,
85                strides.as_ptr(),
86                data_ptr,
87                PySliceContainer::from(self),
88            )
89        }
90    }
91}
92
93impl<A, D> IntoPyArray for ArrayBase<OwnedRepr<A>, D>
94where
95    A: Element,
96    D: Dimension,
97{
98    type Item = A;
99    type Dim = D;
100
101    fn into_pyarray<'py>(self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
102        PyArray::from_owned_array(py, self)
103    }
104}
105
106/// Conversion trait from borrowing Rust types to [`PyArray`].
107///
108/// This trait takes `&self` by reference, which means it allocates in Python heap and then copies the elements there.
109///
110/// # Examples
111///
112/// ```
113/// use numpy::{PyArray, ToPyArray, PyArrayMethods};
114/// use pyo3::Python;
115///
116/// Python::with_gil(|py| {
117///     let py_array = vec![1, 2, 3].to_pyarray(py);
118///
119///     assert_eq!(py_array.readonly().as_slice().unwrap(), &[1, 2, 3]);
120/// });
121/// ```
122///
123/// Due to copying the elments, this method converts non-contiguous arrays to C-order contiguous arrays.
124///
125/// ```
126/// use numpy::prelude::*;
127/// use numpy::{PyArray, ToPyArray};
128/// use ndarray::{arr3, s};
129/// use pyo3::Python;
130///
131/// Python::with_gil(|py| {
132///     let array = arr3(&[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]);
133///     let py_array = array.slice(s![.., 0..1, ..]).to_pyarray(py);
134///
135///     assert_eq!(py_array.readonly().as_array(), arr3(&[[[1, 2, 3]], [[7, 8, 9]]]));
136///     assert!(py_array.is_c_contiguous());
137/// });
138/// ```
139pub trait ToPyArray {
140    /// The element type of resulting array.
141    type Item: Element;
142    /// The dimension type of the resulting array.
143    type Dim: Dimension;
144
145    /// Copies the content pointed to by `&self` into a newly allocated NumPy array.
146    fn to_pyarray<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>>;
147
148    /// Deprecated name for [ToPyArray::to_pyarray`].
149    #[deprecated(since = "0.23.0", note = "renamed to ToPyArray::to_pyarray`")]
150    #[inline]
151    fn to_pyarray_bound<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
152        self.to_pyarray(py)
153    }
154}
155
156impl<T: Element> ToPyArray for [T] {
157    type Item = T;
158    type Dim = Ix1;
159
160    fn to_pyarray<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
161        PyArray::from_slice(py, self)
162    }
163}
164
165impl<S, D, A> ToPyArray for ArrayBase<S, D>
166where
167    S: Data<Elem = A>,
168    D: Dimension,
169    A: Element,
170{
171    type Item = A;
172    type Dim = D;
173
174    fn to_pyarray<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
175        let len = self.len();
176        match self.order() {
177            Some(flag) if A::IS_COPY => {
178                // if the array is contiguous, copy it by `copy_nonoverlapping`.
179                let strides = self.npy_strides();
180                unsafe {
181                    let array = PyArray::new_uninit(py, self.raw_dim(), strides.as_ptr(), flag);
182                    ptr::copy_nonoverlapping(self.as_ptr(), array.data(), len);
183                    array
184                }
185            }
186            _ => {
187                // if the array is not contiguous, copy all elements by `ArrayBase::iter`.
188                let dim = self.raw_dim();
189                unsafe {
190                    let array = PyArray::<A, _>::new(py, dim, false);
191                    let mut data_ptr = array.data();
192                    for item in self.iter() {
193                        data_ptr.write(item.clone_ref(py));
194                        data_ptr = data_ptr.add(1);
195                    }
196                    array
197                }
198            }
199        }
200    }
201}
202
203#[cfg(feature = "nalgebra")]
204impl<N, R, C, S> ToPyArray for nalgebra::Matrix<N, R, C, S>
205where
206    N: nalgebra::Scalar + Element,
207    R: nalgebra::Dim,
208    C: nalgebra::Dim,
209    S: nalgebra::Storage<N, R, C>,
210{
211    type Item = N;
212    type Dim = crate::Ix2;
213
214    /// Note that the NumPy array always has Fortran memory layout
215    /// matching the [memory layout][memory-layout] used by [`nalgebra`].
216    ///
217    /// [memory-layout]: https://nalgebra.org/docs/faq/#what-is-the-memory-layout-of-matrices
218    fn to_pyarray<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
219        unsafe {
220            let array = PyArray::<N, _>::new(py, (self.nrows(), self.ncols()), true);
221            let mut data_ptr = array.data();
222            if self.data.is_contiguous() {
223                ptr::copy_nonoverlapping(self.data.ptr(), data_ptr, self.len());
224            } else {
225                for item in self.iter() {
226                    data_ptr.write(item.clone_ref(py));
227                    data_ptr = data_ptr.add(1);
228                }
229            }
230            array
231        }
232    }
233}
234
235pub(crate) trait ArrayExt {
236    fn npy_strides(&self) -> [npyffi::npy_intp; 32];
237    fn order(&self) -> Option<c_int>;
238}
239
240impl<A, S, D> ArrayExt for ArrayBase<S, D>
241where
242    S: Data<Elem = A>,
243    D: Dimension,
244{
245    fn npy_strides(&self) -> [npyffi::npy_intp; 32] {
246        let strides = self.strides();
247        let itemsize = mem::size_of::<A>() as isize;
248
249        assert!(strides.len() <= 32, "{}", MAX_DIMENSIONALITY_ERR);
250
251        let mut new_strides = [0; 32];
252
253        for i in 0..strides.len() {
254            new_strides[i] = (strides[i] * itemsize) as npyffi::npy_intp;
255        }
256
257        new_strides
258    }
259
260    fn order(&self) -> Option<c_int> {
261        if self.is_standard_layout() {
262            Some(npyffi::NPY_ORDER::NPY_CORDER as _)
263        } else if self.ndim() > 1 && self.raw_view().reversed_axes().is_standard_layout() {
264            Some(npyffi::NPY_ORDER::NPY_FORTRANORDER as _)
265        } else {
266            None
267        }
268    }
269}
270
271/// Utility trait to specify the dimensions of an array.
272pub trait ToNpyDims: Dimension + Sealed {
273    #[doc(hidden)]
274    fn ndim_cint(&self) -> c_int {
275        self.ndim() as c_int
276    }
277    #[doc(hidden)]
278    fn as_dims_ptr(&mut self) -> *mut npyffi::npy_intp {
279        self.slice_mut().as_ptr() as *mut npyffi::npy_intp
280    }
281    #[doc(hidden)]
282    fn to_npy_dims(&mut self) -> npyffi::PyArray_Dims {
283        npyffi::PyArray_Dims {
284            ptr: self.as_dims_ptr(),
285            len: self.ndim_cint(),
286        }
287    }
288}
289
290mod sealed {
291    pub trait Sealed {}
292}
293
294use sealed::Sealed;
295
296impl<D> ToNpyDims for D where D: Dimension {}
297
298/// Trait implemented by types that can be used to index an array.
299///
300/// This is equivalent to [`ndarray::NdIndex`] but accounts for
301/// NumPy strides being in units of bytes instead of elements.
302///
303/// All types which implement [`IntoDimension`] implement this trait as well.
304/// This includes at least
305/// - [tuple](https://doc.rust-lang.org/stable/std/primitive.tuple.html)
306/// - [array](https://doc.rust-lang.org/stable/std/primitive.array.html)
307/// - [slice](https://doc.rust-lang.org/stable/std/primitive.slice.html)
308pub trait NpyIndex: IntoDimension + Sealed {
309    #[doc(hidden)]
310    fn get_checked<T>(self, dims: &[usize], strides: &[isize]) -> Option<isize>;
311    #[doc(hidden)]
312    fn get_unchecked<T>(self, strides: &[isize]) -> isize;
313}
314
315impl<D: IntoDimension> Sealed for D {}
316
317impl<D: IntoDimension> NpyIndex for D {
318    fn get_checked<T>(self, dims: &[usize], strides: &[isize]) -> Option<isize> {
319        let indices = self.into_dimension();
320        let indices = indices.slice();
321
322        if indices.len() != dims.len() {
323            return None;
324        }
325        if indices.iter().zip(dims).any(|(i, d)| i >= d) {
326            return None;
327        }
328
329        Some(get_unchecked_impl::<T>(indices, strides))
330    }
331
332    fn get_unchecked<T>(self, strides: &[isize]) -> isize {
333        let indices = self.into_dimension();
334        let indices = indices.slice();
335        get_unchecked_impl::<T>(indices, strides)
336    }
337}
338
339fn get_unchecked_impl<T>(indices: &[usize], strides: &[isize]) -> isize {
340    let size = mem::size_of::<T>() as isize;
341
342    indices
343        .iter()
344        .zip(strides)
345        .map(|(&i, stride)| stride * i as isize / size)
346        .sum()
347}
numpy/convert.rs

numpy/
convert.rs
