numpy/
array.rs

1//! Safe interface for NumPy's [N-dimensional arrays][ndarray]
2//!
3//! [ndarray]: https://numpy.org/doc/stable/reference/arrays.ndarray.html
4
5use std::{
6    marker::PhantomData,
7    mem,
8    os::raw::{c_int, c_void},
9    ptr, slice,
10};
11
12use ndarray::{
13    Array, ArrayBase, ArrayView, ArrayViewMut, Axis, Data, Dim, Dimension, IntoDimension, Ix0, Ix1,
14    Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, RawArrayView, RawArrayViewMut, RawData, ShapeBuilder,
15    StrideShape,
16};
17use num_traits::AsPrimitive;
18use pyo3::{
19    ffi,
20    types::{DerefToPyAny, PyModule},
21    Bound, CastError, Py, PyAny, PyErr, PyResult, PyTypeCheck, PyTypeInfo, Python,
22};
23
24use crate::borrow::{PyReadonlyArray, PyReadwriteArray};
25use crate::cold;
26use crate::convert::{ArrayExt, IntoPyArray, NpyIndex, ToNpyDims, ToPyArray};
27use crate::dtype::{Element, PyArrayDescrMethods};
28use crate::error::{
29    BorrowError, DimensionalityError, FromVecError, IgnoreError, NotContiguousError, TypeError,
30    DIMENSIONALITY_MISMATCH_ERR, MAX_DIMENSIONALITY_ERR,
31};
32use crate::npyffi::{self, npy_intp, NPY_ORDER, PY_ARRAY_API};
33use crate::slice_container::PySliceContainer;
34use crate::untyped_array::{PyUntypedArray, PyUntypedArrayMethods};
35
36/// A safe, statically-typed wrapper for NumPy's [`ndarray`][ndarray] class.
37///
38/// # Memory location
39///
40/// - Allocated by Rust: Constructed via [`IntoPyArray`] or
41///   [`from_vec`][Self::from_vec] or [`from_owned_array`][Self::from_owned_array].
42///
43/// These methods transfers ownership of the Rust allocation into a suitable Python object
44/// and uses the memory as the internal buffer backing the NumPy array.
45///
46/// Please note that some destructive methods like [`resize`][Self::resize] will fail
47/// when used with this kind of array as NumPy cannot reallocate the internal buffer.
48///
49/// - Allocated by NumPy: Constructed via other methods, like [`ToPyArray`] or
50///   [`from_slice`][Self::from_slice] or [`from_array`][Self::from_array].
51///
52/// These methods allocate memory in Python's private heap via NumPy's API.
53///
54/// In both cases, `PyArray` is managed by Python so it can neither be moved from
55/// nor deallocated manually.
56///
57/// # References
58///
59/// Like [`new`][Self::new], all constructor methods of `PyArray` return a shared reference `&PyArray`
60/// instead of an owned value. This design follows [PyO3's ownership concept][pyo3-memory],
61/// i.e. the return value is GIL-bound owning reference into Python's heap.
62///
63/// # Element type and dimensionality
64///
65/// `PyArray` has two type parametes `T` and `D`.
66/// `T` represents the type of its elements, e.g. [`f32`] or [`PyObject`].
67/// `D` represents its dimensionality, e.g [`Ix2`][type@Ix2] or [`IxDyn`][type@IxDyn].
68///
69/// Element types are Rust types which implement the [`Element`] trait.
70/// Dimensions are represented by the [`ndarray::Dimension`] trait.
71///
72/// Typically, `Ix1, Ix2, ...` are used for fixed dimensionality arrays,
73/// and `IxDyn` is used for dynamic dimensionality arrays. Type aliases
74/// for combining `PyArray` with these types are provided, e.g. [`PyArray1`] or [`PyArrayDyn`].
75///
76/// To specify concrete dimension like `3×4×5`, types which implement the [`ndarray::IntoDimension`]
77/// trait are used. Typically, this means arrays like `[3, 4, 5]` or tuples like `(3, 4, 5)`.
78///
79/// # Example
80///
81/// ```
82/// use numpy::{PyArray, PyArrayMethods};
83/// use ndarray::{array, Array};
84/// use pyo3::Python;
85///
86/// Python::attach(|py| {
87///     let pyarray = PyArray::arange(py, 0., 4., 1.).reshape([2, 2]).unwrap();
88///     let array = array![[3., 4.], [5., 6.]];
89///
90///     assert_eq!(
91///         array.dot(&pyarray.readonly().as_array()),
92///         array![[8., 15.], [12., 23.]]
93///     );
94/// });
95/// ```
96///
97/// [ndarray]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html
98/// [pyo3-memory]: https://pyo3.rs/main/memory.html
99#[repr(transparent)]
100pub struct PyArray<T, D>(PyAny, PhantomData<T>, PhantomData<D>);
101
102/// Zero-dimensional array.
103pub type PyArray0<T> = PyArray<T, Ix0>;
104/// One-dimensional array.
105pub type PyArray1<T> = PyArray<T, Ix1>;
106/// Two-dimensional array.
107pub type PyArray2<T> = PyArray<T, Ix2>;
108/// Three-dimensional array.
109pub type PyArray3<T> = PyArray<T, Ix3>;
110/// Four-dimensional array.
111pub type PyArray4<T> = PyArray<T, Ix4>;
112/// Five-dimensional array.
113pub type PyArray5<T> = PyArray<T, Ix5>;
114/// Six-dimensional array.
115pub type PyArray6<T> = PyArray<T, Ix6>;
116/// Dynamic-dimensional array.
117pub type PyArrayDyn<T> = PyArray<T, IxDyn>;
118
119/// Returns a handle to NumPy's multiarray module.
120pub fn get_array_module<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyModule>> {
121    PyModule::import(py, npyffi::array::mod_name(py)?)
122}
123
124impl<T, D> DerefToPyAny for PyArray<T, D> {}
125
126unsafe impl<T: Element, D: Dimension> PyTypeInfo for PyArray<T, D> {
127    const NAME: &'static str = "PyArray<T, D>";
128    const MODULE: Option<&'static str> = Some("numpy");
129
130    fn type_object_raw<'py>(py: Python<'py>) -> *mut ffi::PyTypeObject {
131        unsafe { npyffi::PY_ARRAY_API.get_type_object(py, npyffi::NpyTypes::PyArray_Type) }
132    }
133
134    fn is_type_of(ob: &Bound<'_, PyAny>) -> bool {
135        Self::extract::<IgnoreError>(ob).is_ok()
136    }
137}
138
139impl<T: Element, D: Dimension> PyArray<T, D> {
140    fn extract<'a, 'py, E>(ob: &'a Bound<'py, PyAny>) -> Result<&'a Bound<'py, Self>, E>
141    where
142        E: From<CastError<'a, 'py>> + From<DimensionalityError> + From<TypeError<'py>>,
143    {
144        // Check if the object is an array.
145        let array = unsafe {
146            if npyffi::PyArray_Check(ob.py(), ob.as_ptr()) == 0 {
147                return Err(CastError::new(
148                    ob.as_borrowed(),
149                    <Self as PyTypeCheck>::classinfo_object(ob.py()),
150                )
151                .into());
152            }
153            ob.cast_unchecked::<Self>()
154        };
155
156        // Check if the dimensionality matches `D`.
157        let src_ndim = array.ndim();
158        if let Some(dst_ndim) = D::NDIM {
159            if src_ndim != dst_ndim {
160                return Err(DimensionalityError::new(src_ndim, dst_ndim).into());
161            }
162        }
163
164        // Check if the element type matches `T`.
165        let src_dtype = array.dtype();
166        let dst_dtype = T::get_dtype(ob.py());
167        if !src_dtype.is_equiv_to(&dst_dtype) {
168            return Err(TypeError::new(src_dtype, dst_dtype).into());
169        }
170
171        Ok(array)
172    }
173
174    /// Creates a new uninitialized NumPy array.
175    ///
176    /// If `is_fortran` is true, then it has Fortran/column-major order,
177    /// otherwise it has C/row-major order.
178    ///
179    /// # Safety
180    ///
181    /// The returned array will always be safe to be dropped as the elements must either
182    /// be trivially copyable (as indicated by `<T as Element>::IS_COPY`) or be pointers
183    /// into Python's heap, which NumPy will automatically zero-initialize.
184    ///
185    /// However, the elements themselves will not be valid and should be initialized manually
186    /// using raw pointers obtained via [`uget_raw`][Self::uget_raw]. Before that, all methods
187    /// which produce references to the elements invoke undefined behaviour. In particular,
188    /// zero-initialized pointers are _not_ valid instances of `PyObject`.
189    ///
190    /// # Example
191    ///
192    /// ```
193    /// use numpy::prelude::*;
194    /// use numpy::PyArray3;
195    /// use pyo3::Python;
196    ///
197    /// Python::attach(|py| {
198    ///     let arr = unsafe {
199    ///         let arr = PyArray3::<i32>::new(py, [4, 5, 6], false);
200    ///
201    ///         for i in 0..4 {
202    ///             for j in 0..5 {
203    ///                 for k in 0..6 {
204    ///                     arr.uget_raw([i, j, k]).write((i * j * k) as i32);
205    ///                 }
206    ///             }
207    ///         }
208    ///
209    ///         arr
210    ///     };
211    ///
212    ///     assert_eq!(arr.shape(), &[4, 5, 6]);
213    /// });
214    /// ```
215    pub unsafe fn new<'py, ID>(py: Python<'py>, dims: ID, is_fortran: bool) -> Bound<'py, Self>
216    where
217        ID: IntoDimension<Dim = D>,
218    {
219        let flags = c_int::from(is_fortran);
220        Self::new_uninit(py, dims, ptr::null_mut(), flags)
221    }
222
223    pub(crate) unsafe fn new_uninit<'py, ID>(
224        py: Python<'py>,
225        dims: ID,
226        strides: *const npy_intp,
227        flag: c_int,
228    ) -> Bound<'py, Self>
229    where
230        ID: IntoDimension<Dim = D>,
231    {
232        let mut dims = dims.into_dimension();
233        let ptr = PY_ARRAY_API.PyArray_NewFromDescr(
234            py,
235            PY_ARRAY_API.get_type_object(py, npyffi::NpyTypes::PyArray_Type),
236            T::get_dtype(py).into_dtype_ptr(),
237            dims.ndim_cint(),
238            dims.as_dims_ptr(),
239            strides as *mut npy_intp, // strides
240            ptr::null_mut(),          // data
241            flag,                     // flag
242            ptr::null_mut(),          // obj
243        );
244
245        Bound::from_owned_ptr(py, ptr).cast_into_unchecked()
246    }
247
248    unsafe fn new_with_data<'py, ID>(
249        py: Python<'py>,
250        dims: ID,
251        strides: *const npy_intp,
252        data_ptr: *const T,
253        container: *mut PyAny,
254    ) -> Bound<'py, Self>
255    where
256        ID: IntoDimension<Dim = D>,
257    {
258        let mut dims = dims.into_dimension();
259        let ptr = PY_ARRAY_API.PyArray_NewFromDescr(
260            py,
261            PY_ARRAY_API.get_type_object(py, npyffi::NpyTypes::PyArray_Type),
262            T::get_dtype(py).into_dtype_ptr(),
263            dims.ndim_cint(),
264            dims.as_dims_ptr(),
265            strides as *mut npy_intp,    // strides
266            data_ptr as *mut c_void,     // data
267            npyffi::NPY_ARRAY_WRITEABLE, // flag
268            ptr::null_mut(),             // obj
269        );
270
271        PY_ARRAY_API.PyArray_SetBaseObject(
272            py,
273            ptr as *mut npyffi::PyArrayObject,
274            container as *mut ffi::PyObject,
275        );
276
277        Bound::from_owned_ptr(py, ptr).cast_into_unchecked()
278    }
279
280    pub(crate) unsafe fn from_raw_parts<'py>(
281        py: Python<'py>,
282        dims: D,
283        strides: *const npy_intp,
284        data_ptr: *const T,
285        container: PySliceContainer,
286    ) -> Bound<'py, Self> {
287        let container = Bound::new(py, container)
288            .expect("Failed to create slice container")
289            .into_ptr();
290
291        Self::new_with_data(py, dims, strides, data_ptr, container.cast())
292    }
293
294    /// Creates a NumPy array backed by `array` and ties its ownership to the Python object `container`.
295    ///
296    /// The resulting NumPy array will be writeable from Python space.  If this is undesireable, use
297    /// [PyReadwriteArray::make_nonwriteable].
298    ///
299    /// # Safety
300    ///
301    /// `container` is set as a base object of the returned array which must not be dropped until `container` is dropped.
302    /// Furthermore, `array` must not be reallocated from the time this method is called and until `container` is dropped.
303    ///
304    /// # Example
305    ///
306    /// ```rust
307    /// # use pyo3::prelude::*;
308    /// # use numpy::{ndarray::Array1, PyArray1};
309    /// #
310    /// #[pyclass]
311    /// struct Owner {
312    ///     array: Array1<f64>,
313    /// }
314    ///
315    /// #[pymethods]
316    /// impl Owner {
317    ///     #[getter]
318    ///     fn array<'py>(this: Bound<'py, Self>) -> Bound<'py, PyArray1<f64>> {
319    ///         let array = &this.borrow().array;
320    ///
321    ///         // SAFETY: The memory backing `array` will stay valid as long as this object is alive
322    ///         // as we do not modify `array` in any way which would cause it to be reallocated.
323    ///         unsafe { PyArray1::borrow_from_array(array, this.into_any()) }
324    ///     }
325    /// }
326    /// ```
327    pub unsafe fn borrow_from_array<'py, S>(
328        array: &ArrayBase<S, D>,
329        container: Bound<'py, PyAny>,
330    ) -> Bound<'py, Self>
331    where
332        S: Data<Elem = T>,
333    {
334        let (strides, dims) = (array.npy_strides(), array.raw_dim());
335        let data_ptr = array.as_ptr();
336
337        let py = container.py();
338
339        Self::new_with_data(
340            py,
341            dims,
342            strides.as_ptr(),
343            data_ptr,
344            container.into_ptr().cast(),
345        )
346    }
347
348    /// Construct a new NumPy array filled with zeros.
349    ///
350    /// If `is_fortran` is true, then it has Fortran/column-major order,
351    /// otherwise it has C/row-major order.
352    ///
353    /// For arrays of Python objects, this will fill the array
354    /// with valid pointers to zero-valued Python integer objects.
355    ///
356    /// See also [`numpy.zeros`][numpy-zeros] and [`PyArray_Zeros`][PyArray_Zeros].
357    ///
358    /// # Example
359    ///
360    /// ```
361    /// use numpy::{PyArray2, PyArrayMethods};
362    /// use pyo3::Python;
363    ///
364    /// Python::attach(|py| {
365    ///     let pyarray = PyArray2::<usize>::zeros(py, [2, 2], true);
366    ///
367    ///     assert_eq!(pyarray.readonly().as_slice().unwrap(), [0; 4]);
368    /// });
369    /// ```
370    ///
371    /// [numpy-zeros]: https://numpy.org/doc/stable/reference/generated/numpy.zeros.html
372    /// [PyArray_Zeros]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_Zeros
373    pub fn zeros<ID>(py: Python<'_>, dims: ID, is_fortran: bool) -> Bound<'_, Self>
374    where
375        ID: IntoDimension<Dim = D>,
376    {
377        let mut dims = dims.into_dimension();
378        unsafe {
379            let ptr = PY_ARRAY_API.PyArray_Zeros(
380                py,
381                dims.ndim_cint(),
382                dims.as_dims_ptr(),
383                T::get_dtype(py).into_dtype_ptr(),
384                if is_fortran { -1 } else { 0 },
385            );
386            Bound::from_owned_ptr(py, ptr).cast_into_unchecked()
387        }
388    }
389
390    /// Constructs a NumPy from an [`ndarray::Array`]
391    ///
392    /// This method uses the internal [`Vec`] of the [`ndarray::Array`] as the base object of the NumPy array.
393    ///
394    /// # Example
395    ///
396    /// ```
397    /// use numpy::{PyArray, PyArrayMethods};
398    /// use ndarray::array;
399    /// use pyo3::Python;
400    ///
401    /// Python::attach(|py| {
402    ///     let pyarray = PyArray::from_owned_array(py, array![[1, 2], [3, 4]]);
403    ///
404    ///     assert_eq!(pyarray.readonly().as_array(), array![[1, 2], [3, 4]]);
405    /// });
406    /// ```
407    pub fn from_owned_array(py: Python<'_>, mut arr: Array<T, D>) -> Bound<'_, Self> {
408        let (strides, dims) = (arr.npy_strides(), arr.raw_dim());
409        let data_ptr = arr.as_mut_ptr();
410        unsafe {
411            Self::from_raw_parts(
412                py,
413                dims,
414                strides.as_ptr(),
415                data_ptr,
416                PySliceContainer::from(arr),
417            )
418        }
419    }
420
421    /// Construct a NumPy array from a [`ndarray::ArrayBase`].
422    ///
423    /// This method allocates memory in Python's heap via the NumPy API,
424    /// and then copies all elements of the array there.
425    ///
426    /// # Example
427    ///
428    /// ```
429    /// use numpy::{PyArray, PyArrayMethods};
430    /// use ndarray::array;
431    /// use pyo3::Python;
432    ///
433    /// Python::attach(|py| {
434    ///     let pyarray = PyArray::from_array(py, &array![[1, 2], [3, 4]]);
435    ///
436    ///     assert_eq!(pyarray.readonly().as_array(), array![[1, 2], [3, 4]]);
437    /// });
438    /// ```
439    pub fn from_array<'py, S>(py: Python<'py>, arr: &ArrayBase<S, D>) -> Bound<'py, Self>
440    where
441        S: Data<Elem = T>,
442    {
443        ToPyArray::to_pyarray(arr, py)
444    }
445}
446
447impl<D: Dimension> PyArray<Py<PyAny>, D> {
448    /// Construct a NumPy array containing objects stored in a [`ndarray::Array`]
449    ///
450    /// This method uses the internal [`Vec`] of the [`ndarray::Array`] as the base object of the NumPy array.
451    ///
452    /// # Example
453    ///
454    /// ```
455    /// use ndarray::array;
456    /// use pyo3::{pyclass, Py, Python, types::PyAnyMethods};
457    /// use numpy::{PyArray, PyArrayMethods};
458    ///
459    /// #[pyclass]
460    /// # #[allow(dead_code)]
461    /// struct CustomElement {
462    ///     foo: i32,
463    ///     bar: f64,
464    /// }
465    ///
466    /// Python::attach(|py| {
467    ///     let array = array![
468    ///         Py::new(py, CustomElement {
469    ///             foo: 1,
470    ///             bar: 2.0,
471    ///         }).unwrap(),
472    ///         Py::new(py, CustomElement {
473    ///             foo: 3,
474    ///             bar: 4.0,
475    ///         }).unwrap(),
476    ///     ];
477    ///
478    ///     let pyarray = PyArray::from_owned_object_array(py, array);
479    ///
480    ///     assert!(pyarray.readonly().as_array().get(0).unwrap().bind(py).is_instance_of::<CustomElement>());
481    /// });
482    /// ```
483    pub fn from_owned_object_array<T>(py: Python<'_>, mut arr: Array<Py<T>, D>) -> Bound<'_, Self> {
484        let (strides, dims) = (arr.npy_strides(), arr.raw_dim());
485        let data_ptr = arr.as_mut_ptr().cast::<Py<PyAny>>().cast_const();
486        unsafe {
487            Self::from_raw_parts(
488                py,
489                dims,
490                strides.as_ptr(),
491                data_ptr,
492                PySliceContainer::from(arr),
493            )
494        }
495    }
496}
497
498impl<T: Element> PyArray<T, Ix1> {
499    /// Construct a one-dimensional array from a [mod@slice].
500    ///
501    /// # Example
502    ///
503    /// ```
504    /// use numpy::{PyArray, PyArrayMethods};
505    /// use pyo3::Python;
506    ///
507    /// Python::attach(|py| {
508    ///     let slice = &[1, 2, 3, 4, 5];
509    ///     let pyarray = PyArray::from_slice(py, slice);
510    ///     assert_eq!(pyarray.readonly().as_slice().unwrap(), &[1, 2, 3, 4, 5]);
511    /// });
512    /// ```
513    pub fn from_slice<'py>(py: Python<'py>, slice: &[T]) -> Bound<'py, Self> {
514        unsafe {
515            let array = PyArray::new(py, [slice.len()], false);
516            let mut data_ptr = array.data();
517            clone_elements(py, slice, &mut data_ptr);
518            array
519        }
520    }
521
522    /// Construct a one-dimensional array from a [`Vec<T>`][Vec].
523    ///
524    /// # Example
525    ///
526    /// ```
527    /// use numpy::{PyArray, PyArrayMethods};
528    /// use pyo3::Python;
529    ///
530    /// Python::attach(|py| {
531    ///     let vec = vec![1, 2, 3, 4, 5];
532    ///     let pyarray = PyArray::from_vec(py, vec);
533    ///     assert_eq!(pyarray.readonly().as_slice().unwrap(), &[1, 2, 3, 4, 5]);
534    /// });
535    /// ```
536    #[inline(always)]
537    pub fn from_vec<'py>(py: Python<'py>, vec: Vec<T>) -> Bound<'py, Self> {
538        vec.into_pyarray(py)
539    }
540
541    /// Construct a one-dimensional array from an [`Iterator`].
542    ///
543    /// If no reliable [`size_hint`][Iterator::size_hint] is available,
544    /// this method can allocate memory multiple times, which can hurt performance.
545    ///
546    /// # Example
547    ///
548    /// ```
549    /// use numpy::{PyArray, PyArrayMethods};
550    /// use pyo3::Python;
551    ///
552    /// Python::attach(|py| {
553    ///     let pyarray = PyArray::from_iter(py, "abcde".chars().map(u32::from));
554    ///     assert_eq!(pyarray.readonly().as_slice().unwrap(), &[97, 98, 99, 100, 101]);
555    /// });
556    /// ```
557    pub fn from_iter<I>(py: Python<'_>, iter: I) -> Bound<'_, Self>
558    where
559        I: IntoIterator<Item = T>,
560    {
561        let data = iter.into_iter().collect::<Vec<_>>();
562        data.into_pyarray(py)
563    }
564}
565
566impl<T: Element> PyArray<T, Ix2> {
567    /// Construct a two-dimension array from a [`Vec<Vec<T>>`][Vec].
568    ///
569    /// This function checks all dimensions of the inner vectors and returns
570    /// an error if they are not all equal.
571    ///
572    /// # Example
573    ///
574    /// ```
575    /// use numpy::{PyArray, PyArrayMethods};
576    /// use pyo3::Python;
577    /// use ndarray::array;
578    ///
579    /// Python::attach(|py| {
580    ///     let vec2 = vec![vec![11, 12], vec![21, 22]];
581    ///     let pyarray = PyArray::from_vec2(py, &vec2).unwrap();
582    ///     assert_eq!(pyarray.readonly().as_array(), array![[11, 12], [21, 22]]);
583    ///
584    ///     let ragged_vec2 = vec![vec![11, 12], vec![21]];
585    ///     assert!(PyArray::from_vec2(py, &ragged_vec2).is_err());
586    /// });
587    /// ```
588    pub fn from_vec2<'py>(py: Python<'py>, v: &[Vec<T>]) -> Result<Bound<'py, Self>, FromVecError> {
589        let len2 = v.first().map_or(0, |v| v.len());
590        let dims = [v.len(), len2];
591        // SAFETY: The result of `Self::new` is always safe to drop.
592        unsafe {
593            let array = Self::new(py, dims, false);
594            let mut data_ptr = array.data();
595            for v in v {
596                if v.len() != len2 {
597                    cold();
598                    return Err(FromVecError::new(v.len(), len2));
599                }
600                clone_elements(py, v, &mut data_ptr);
601            }
602            Ok(array)
603        }
604    }
605}
606
607impl<T: Element> PyArray<T, Ix3> {
608    /// Construct a three-dimensional array from a [`Vec<Vec<Vec<T>>>`][Vec].
609    ///
610    /// This function checks all dimensions of the inner vectors and returns
611    /// an error if they are not all equal.
612    ///
613    /// # Example
614    ///
615    /// ```
616    /// use numpy::{PyArray, PyArrayMethods};
617    /// use pyo3::Python;
618    /// use ndarray::array;
619    ///
620    /// Python::attach(|py| {
621    ///     let vec3 = vec![
622    ///         vec![vec![111, 112], vec![121, 122]],
623    ///         vec![vec![211, 212], vec![221, 222]],
624    ///     ];
625    ///     let pyarray = PyArray::from_vec3(py, &vec3).unwrap();
626    ///     assert_eq!(
627    ///         pyarray.readonly().as_array(),
628    ///         array![[[111, 112], [121, 122]], [[211, 212], [221, 222]]]
629    ///     );
630    ///
631    ///     let ragged_vec3 = vec![
632    ///         vec![vec![111, 112], vec![121, 122]],
633    ///         vec![vec![211], vec![221, 222]],
634    ///     ];
635    ///     assert!(PyArray::from_vec3(py, &ragged_vec3).is_err());
636    /// });
637    /// ```
638    pub fn from_vec3<'py>(
639        py: Python<'py>,
640        v: &[Vec<Vec<T>>],
641    ) -> Result<Bound<'py, Self>, FromVecError> {
642        let len2 = v.first().map_or(0, |v| v.len());
643        let len3 = v.first().map_or(0, |v| v.first().map_or(0, |v| v.len()));
644        let dims = [v.len(), len2, len3];
645        // SAFETY: The result of `Self::new` is always safe to drop.
646        unsafe {
647            let array = Self::new(py, dims, false);
648            let mut data_ptr = array.data();
649            for v in v {
650                if v.len() != len2 {
651                    cold();
652                    return Err(FromVecError::new(v.len(), len2));
653                }
654                for v in v {
655                    if v.len() != len3 {
656                        cold();
657                        return Err(FromVecError::new(v.len(), len3));
658                    }
659                    clone_elements(py, v, &mut data_ptr);
660                }
661            }
662            Ok(array)
663        }
664    }
665}
666
667impl<T: Element + AsPrimitive<f64>> PyArray<T, Ix1> {
668    /// Return evenly spaced values within a given interval.
669    ///
670    /// See [numpy.arange][numpy.arange] for the Python API and [PyArray_Arange][PyArray_Arange] for the C API.
671    ///
672    /// # Example
673    ///
674    /// ```
675    /// use numpy::{PyArray, PyArrayMethods};
676    /// use pyo3::Python;
677    ///
678    /// Python::attach(|py| {
679    ///     let pyarray = PyArray::arange(py, 2.0, 4.0, 0.5);
680    ///     assert_eq!(pyarray.readonly().as_slice().unwrap(), &[2.0, 2.5, 3.0, 3.5]);
681    ///
682    ///     let pyarray = PyArray::arange(py, -2, 4, 3);
683    ///     assert_eq!(pyarray.readonly().as_slice().unwrap(), &[-2, 1]);
684    /// });
685    /// ```
686    ///
687    /// [numpy.arange]: https://numpy.org/doc/stable/reference/generated/numpy.arange.html
688    /// [PyArray_Arange]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_Arange
689    pub fn arange<'py>(py: Python<'py>, start: T, stop: T, step: T) -> Bound<'py, Self> {
690        unsafe {
691            let ptr = PY_ARRAY_API.PyArray_Arange(
692                py,
693                start.as_(),
694                stop.as_(),
695                step.as_(),
696                T::get_dtype(py).num(),
697            );
698            Bound::from_owned_ptr(py, ptr).cast_into_unchecked()
699        }
700    }
701}
702
703unsafe fn clone_elements<T: Element>(py: Python<'_>, elems: &[T], data_ptr: &mut *mut T) {
704    if T::IS_COPY {
705        ptr::copy_nonoverlapping(elems.as_ptr(), *data_ptr, elems.len());
706        *data_ptr = data_ptr.add(elems.len());
707    } else {
708        for elem in elems {
709            data_ptr.write(elem.clone_ref(py));
710            *data_ptr = data_ptr.add(1);
711        }
712    }
713}
714
715/// Implementation of functionality for [`PyArray<T, D>`].
716#[doc(alias = "PyArray")]
717pub trait PyArrayMethods<'py, T, D>: PyUntypedArrayMethods<'py> {
718    /// Access an untyped representation of this array.
719    fn as_untyped(&self) -> &Bound<'py, PyUntypedArray>;
720
721    /// Returns a pointer to the first element of the array.
722    fn data(&self) -> *mut T;
723
724    /// Same as [`shape`][PyUntypedArray::shape], but returns `D` instead of `&[usize]`.
725    #[inline(always)]
726    fn dims(&self) -> D
727    where
728        D: Dimension,
729    {
730        D::from_dimension(&Dim(self.shape())).expect(DIMENSIONALITY_MISMATCH_ERR)
731    }
732
733    /// Returns an immutable view of the internal data as a slice.
734    ///
735    /// # Safety
736    ///
737    /// Calling this method is undefined behaviour if the underlying array
738    /// is aliased mutably by other instances of `PyArray`
739    /// or concurrently modified by Python or other native code.
740    ///
741    /// Please consider the safe alternative [`PyReadonlyArray::as_slice`].
742    unsafe fn as_slice(&self) -> Result<&[T], NotContiguousError>
743    where
744        T: Element,
745        D: Dimension,
746    {
747        if self.is_contiguous() {
748            Ok(slice::from_raw_parts(self.data(), self.len()))
749        } else {
750            Err(NotContiguousError)
751        }
752    }
753
754    /// Returns a mutable view of the internal data as a slice.
755    ///
756    /// # Safety
757    ///
758    /// Calling this method is undefined behaviour if the underlying array
759    /// is aliased immutably or mutably by other instances of [`PyArray`]
760    /// or concurrently modified by Python or other native code.
761    ///
762    /// Please consider the safe alternative [`PyReadwriteArray::as_slice_mut`].
763    #[allow(clippy::mut_from_ref)]
764    unsafe fn as_slice_mut(&self) -> Result<&mut [T], NotContiguousError>
765    where
766        T: Element,
767        D: Dimension,
768    {
769        if self.is_contiguous() {
770            Ok(slice::from_raw_parts_mut(self.data(), self.len()))
771        } else {
772            Err(NotContiguousError)
773        }
774    }
775
776    /// Get a reference of the specified element if the given index is valid.
777    ///
778    /// # Safety
779    ///
780    /// Calling this method is undefined behaviour if the underlying array
781    /// is aliased mutably by other instances of `PyArray`
782    /// or concurrently modified by Python or other native code.
783    ///
784    /// Consider using safe alternatives like [`PyReadonlyArray::get`].
785    ///
786    /// # Example
787    ///
788    /// ```
789    /// use numpy::{PyArray, PyArrayMethods};
790    /// use pyo3::Python;
791    ///
792    /// Python::attach(|py| {
793    ///     let pyarray = PyArray::arange(py, 0, 16, 1).reshape([2, 2, 4]).unwrap();
794    ///
795    ///     assert_eq!(unsafe { *pyarray.get([1, 0, 3]).unwrap() }, 11);
796    /// });
797    /// ```
798    unsafe fn get(&self, index: impl NpyIndex<Dim = D>) -> Option<&T>
799    where
800        T: Element,
801        D: Dimension;
802
803    /// Same as [`get`][Self::get], but returns `Option<&mut T>`.
804    ///
805    /// # Safety
806    ///
807    /// Calling this method is undefined behaviour if the underlying array
808    /// is aliased immutably or mutably by other instances of [`PyArray`]
809    /// or concurrently modified by Python or other native code.
810    ///
811    /// Consider using safe alternatives like [`PyReadwriteArray::get_mut`].
812    ///
813    /// # Example
814    ///
815    /// ```
816    /// use numpy::{PyArray, PyArrayMethods};
817    /// use pyo3::Python;
818    ///
819    /// Python::attach(|py| {
820    ///     let pyarray = PyArray::arange(py, 0, 16, 1).reshape([2, 2, 4]).unwrap();
821    ///
822    ///     unsafe {
823    ///         *pyarray.get_mut([1, 0, 3]).unwrap() = 42;
824    ///     }
825    ///
826    ///     assert_eq!(unsafe { *pyarray.get([1, 0, 3]).unwrap() }, 42);
827    /// });
828    /// ```
829    #[allow(clippy::mut_from_ref)]
830    unsafe fn get_mut(&self, index: impl NpyIndex<Dim = D>) -> Option<&mut T>
831    where
832        T: Element,
833        D: Dimension;
834
835    /// Get an immutable reference of the specified element,
836    /// without checking the given index.
837    ///
838    /// See [`NpyIndex`] for what types can be used as the index.
839    ///
840    /// # Safety
841    ///
842    /// Passing an invalid index is undefined behavior.
843    /// The element must also have been initialized and
844    /// all other references to it is must also be shared.
845    ///
846    /// See [`PyReadonlyArray::get`] for a safe alternative.
847    ///
848    /// # Example
849    ///
850    /// ```
851    /// use numpy::{PyArray, PyArrayMethods};
852    /// use pyo3::Python;
853    ///
854    /// Python::attach(|py| {
855    ///     let pyarray = PyArray::arange(py, 0, 16, 1).reshape([2, 2, 4]).unwrap();
856    ///
857    ///     assert_eq!(unsafe { *pyarray.uget([1, 0, 3]) }, 11);
858    /// });
859    /// ```
860    #[inline(always)]
861    unsafe fn uget<Idx>(&self, index: Idx) -> &T
862    where
863        T: Element,
864        D: Dimension,
865        Idx: NpyIndex<Dim = D>,
866    {
867        &*self.uget_raw(index)
868    }
869
870    /// Same as [`uget`](Self::uget), but returns `&mut T`.
871    ///
872    /// # Safety
873    ///
874    /// Passing an invalid index is undefined behavior.
875    /// The element must also have been initialized and
876    /// other references to it must not exist.
877    ///
878    /// See [`PyReadwriteArray::get_mut`] for a safe alternative.
879    #[inline(always)]
880    #[allow(clippy::mut_from_ref)]
881    unsafe fn uget_mut<Idx>(&self, index: Idx) -> &mut T
882    where
883        T: Element,
884        D: Dimension,
885        Idx: NpyIndex<Dim = D>,
886    {
887        &mut *self.uget_raw(index)
888    }
889
890    /// Same as [`uget`][Self::uget], but returns `*mut T`.
891    ///
892    /// # Safety
893    ///
894    /// Passing an invalid index is undefined behavior.
895    #[inline(always)]
896    unsafe fn uget_raw<Idx>(&self, index: Idx) -> *mut T
897    where
898        T: Element,
899        D: Dimension,
900        Idx: NpyIndex<Dim = D>,
901    {
902        let offset = index.get_unchecked::<T>(self.strides());
903        self.data().offset(offset) as *mut _
904    }
905
906    /// Get a copy of the specified element in the array.
907    ///
908    /// See [`NpyIndex`] for what types can be used as the index.
909    ///
910    /// # Example
911    /// ```
912    /// use numpy::{PyArray, PyArrayMethods};
913    /// use pyo3::Python;
914    ///
915    /// Python::attach(|py| {
916    ///     let pyarray = PyArray::arange(py, 0, 16, 1).reshape([2, 2, 4]).unwrap();
917    ///
918    ///     assert_eq!(pyarray.get_owned([1, 0, 3]), Some(11));
919    /// });
920    /// ```
921    fn get_owned<Idx>(&self, index: Idx) -> Option<T>
922    where
923        T: Element,
924        D: Dimension,
925        Idx: NpyIndex<Dim = D>;
926
927    /// Turn an array with fixed dimensionality into one with dynamic dimensionality.
928    fn to_dyn(&self) -> &Bound<'py, PyArray<T, IxDyn>>
929    where
930        T: Element,
931        D: Dimension;
932
933    /// Returns a copy of the internal data of the array as a [`Vec`].
934    ///
935    /// Fails if the internal array is not contiguous. See also [`as_slice`][Self::as_slice].
936    ///
937    /// # Example
938    ///
939    /// ```
940    /// use numpy::{PyArray2, PyArrayMethods};
941    /// use pyo3::{Python, types::PyAnyMethods, ffi::c_str};
942    ///
943    /// # fn main() -> pyo3::PyResult<()> {
944    /// Python::attach(|py| {
945    ///     let pyarray= py
946    ///         .eval(c_str!("__import__('numpy').array([[0, 1], [2, 3]], dtype='int64')"), None, None)?
947    ///         .cast_into::<PyArray2<i64>>()?;
948    ///
949    ///     assert_eq!(pyarray.to_vec()?, vec![0, 1, 2, 3]);
950    /// #   Ok(())
951    /// })
952    /// # }
953    /// ```
954    fn to_vec(&self) -> Result<Vec<T>, NotContiguousError>
955    where
956        T: Element,
957        D: Dimension;
958
959    /// Get an immutable borrow of the NumPy array
960    fn try_readonly(&self) -> Result<PyReadonlyArray<'py, T, D>, BorrowError>
961    where
962        T: Element,
963        D: Dimension;
964
965    /// Get an immutable borrow of the NumPy array
966    ///
967    /// # Panics
968    ///
969    /// Panics if the allocation backing the array is currently mutably borrowed.
970    ///
971    /// For a non-panicking variant, use [`try_readonly`][Self::try_readonly].
972    fn readonly(&self) -> PyReadonlyArray<'py, T, D>
973    where
974        T: Element,
975        D: Dimension,
976    {
977        self.try_readonly().unwrap()
978    }
979
980    /// Get a mutable borrow of the NumPy array
981    fn try_readwrite(&self) -> Result<PyReadwriteArray<'py, T, D>, BorrowError>
982    where
983        T: Element,
984        D: Dimension;
985
986    /// Get a mutable borrow of the NumPy array
987    ///
988    /// # Panics
989    ///
990    /// Panics if the allocation backing the array is currently borrowed or
991    /// if the array is [flagged as][flags] not writeable.
992    ///
993    /// For a non-panicking variant, use [`try_readwrite`][Self::try_readwrite].
994    ///
995    /// [flags]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.flags.html
996    fn readwrite(&self) -> PyReadwriteArray<'py, T, D>
997    where
998        T: Element,
999        D: Dimension,
1000    {
1001        self.try_readwrite().unwrap()
1002    }
1003
1004    /// Returns an [`ArrayView`] of the internal array.
1005    ///
1006    /// See also [`PyReadonlyArray::as_array`].
1007    ///
1008    /// # Safety
1009    ///
1010    /// Calling this method invalidates all exclusive references to the internal data, e.g. `&mut [T]` or `ArrayViewMut`.
1011    unsafe fn as_array(&self) -> ArrayView<'_, T, D>
1012    where
1013        T: Element,
1014        D: Dimension;
1015
1016    /// Returns an [`ArrayViewMut`] of the internal array.
1017    ///
1018    /// See also [`PyReadwriteArray::as_array_mut`].
1019    ///
1020    /// # Safety
1021    ///
1022    /// Calling this method invalidates all other references to the internal data, e.g. `ArrayView` or `ArrayViewMut`.
1023    unsafe fn as_array_mut(&self) -> ArrayViewMut<'_, T, D>
1024    where
1025        T: Element,
1026        D: Dimension;
1027
1028    /// Returns the internal array as [`RawArrayView`] enabling element access via raw pointers
1029    fn as_raw_array(&self) -> RawArrayView<T, D>
1030    where
1031        T: Element,
1032        D: Dimension;
1033
1034    /// Returns the internal array as [`RawArrayViewMut`] enabling element access via raw pointers
1035    fn as_raw_array_mut(&self) -> RawArrayViewMut<T, D>
1036    where
1037        T: Element,
1038        D: Dimension;
1039
1040    /// Get a copy of the array as an [`ndarray::Array`].
1041    ///
1042    /// # Example
1043    ///
1044    /// ```
1045    /// use numpy::{PyArray, PyArrayMethods};
1046    /// use ndarray::array;
1047    /// use pyo3::Python;
1048    ///
1049    /// Python::attach(|py| {
1050    ///     let pyarray = PyArray::arange(py, 0, 4, 1).reshape([2, 2]).unwrap();
1051    ///
1052    ///     assert_eq!(
1053    ///         pyarray.to_owned_array(),
1054    ///         array![[0, 1], [2, 3]]
1055    ///     )
1056    /// });
1057    /// ```
1058    fn to_owned_array(&self) -> Array<T, D>
1059    where
1060        T: Element,
1061        D: Dimension;
1062
1063    /// Copies `self` into `other`, performing a data type conversion if necessary.
1064    ///
1065    /// See also [`PyArray_CopyInto`][PyArray_CopyInto].
1066    ///
1067    /// # Example
1068    ///
1069    /// ```
1070    /// use numpy::{PyArray, PyArrayMethods};
1071    /// use pyo3::Python;
1072    ///
1073    /// Python::attach(|py| {
1074    ///     let pyarray_f = PyArray::arange(py, 2.0, 5.0, 1.0);
1075    ///     let pyarray_i = unsafe { PyArray::<i64, _>::new(py, [3], false) };
1076    ///
1077    ///     assert!(pyarray_f.copy_to(&pyarray_i).is_ok());
1078    ///
1079    ///     assert_eq!(pyarray_i.readonly().as_slice().unwrap(), &[2, 3, 4]);
1080    /// });
1081    /// ```
1082    ///
1083    /// [PyArray_CopyInto]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_CopyInto
1084    fn copy_to<U: Element>(&self, other: &Bound<'py, PyArray<U, D>>) -> PyResult<()>
1085    where
1086        T: Element;
1087
1088    /// Deprecated version of [`cast_array`](PyArrayMethods::cast_array)
1089    #[deprecated(since = "0.26.0", note = "use `cast_array` instead")]
1090    #[inline]
1091    fn cast<U: Element>(&self, is_fortran: bool) -> PyResult<Bound<'py, PyArray<U, D>>>
1092    where
1093        T: Element,
1094    {
1095        self.cast_array(is_fortran)
1096    }
1097
1098    /// Cast the `PyArray<T>` to `PyArray<U>`, by allocating a new array.
1099    ///
1100    /// See also [`PyArray_CastToType`][PyArray_CastToType].
1101    ///
1102    /// # Example
1103    ///
1104    /// ```
1105    /// use numpy::{PyArray, PyArrayMethods};
1106    /// use pyo3::Python;
1107    ///
1108    /// Python::attach(|py| {
1109    ///     let pyarray_f = PyArray::arange(py, 2.0, 5.0, 1.0);
1110    ///
1111    ///     let pyarray_i = pyarray_f.cast_array::<i32>(false).unwrap();
1112    ///
1113    ///     assert_eq!(pyarray_i.readonly().as_slice().unwrap(), &[2, 3, 4]);
1114    /// });
1115    /// ```
1116    ///
1117    /// [PyArray_CastToType]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_CastToType
1118    fn cast_array<U: Element>(&self, is_fortran: bool) -> PyResult<Bound<'py, PyArray<U, D>>>
1119    where
1120        T: Element;
1121
1122    /// A view of `self` with a different order of axes determined by `axes`.
1123    ///
1124    /// If `axes` is `None`, the order of axes is reversed which corresponds to the standard matrix transpose.
1125    ///
1126    /// See also [`numpy.transpose`][numpy-transpose] and [`PyArray_Transpose`][PyArray_Transpose].
1127    ///
1128    /// # Example
1129    ///
1130    /// ```
1131    /// use numpy::prelude::*;
1132    /// use numpy::PyArray;
1133    /// use pyo3::Python;
1134    /// use ndarray::array;
1135    ///
1136    /// Python::attach(|py| {
1137    ///     let array = array![[0, 1, 2], [3, 4, 5]].into_pyarray(py);
1138    ///
1139    ///     let array = array.permute(Some([1, 0])).unwrap();
1140    ///
1141    ///     assert_eq!(array.readonly().as_array(), array![[0, 3], [1, 4], [2, 5]]);
1142    /// });
1143    /// ```
1144    ///
1145    /// [numpy-transpose]: https://numpy.org/doc/stable/reference/generated/numpy.transpose.html
1146    /// [PyArray_Transpose]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_Transpose
1147    fn permute<ID: IntoDimension>(&self, axes: Option<ID>) -> PyResult<Bound<'py, PyArray<T, D>>>
1148    where
1149        T: Element;
1150
1151    /// Special case of [`permute`][Self::permute] which reverses the order the axes.
1152    fn transpose(&self) -> PyResult<Bound<'py, PyArray<T, D>>>
1153    where
1154        T: Element,
1155    {
1156        self.permute::<()>(None)
1157    }
1158
1159    /// Construct a new array which has same values as `self`,
1160    /// but has different dimensions specified by `shape`
1161    /// and a possibly different memory order specified by `order`.
1162    ///
1163    /// See also [`numpy.reshape`][numpy-reshape] and [`PyArray_Newshape`][PyArray_Newshape].
1164    ///
1165    /// # Example
1166    ///
1167    /// ```
1168    /// use numpy::prelude::*;
1169    /// use numpy::{npyffi::NPY_ORDER, PyArray};
1170    /// use pyo3::Python;
1171    /// use ndarray::array;
1172    ///
1173    /// Python::attach(|py| {
1174    ///     let array =
1175    ///         PyArray::from_iter(py, 0..9).reshape_with_order([3, 3], NPY_ORDER::NPY_FORTRANORDER).unwrap();
1176    ///
1177    ///     assert_eq!(array.readonly().as_array(), array![[0, 3, 6], [1, 4, 7], [2, 5, 8]]);
1178    ///     assert!(array.is_fortran_contiguous());
1179    ///
1180    ///     assert!(array.reshape([5]).is_err());
1181    /// });
1182    /// ```
1183    ///
1184    /// [numpy-reshape]: https://numpy.org/doc/stable/reference/generated/numpy.reshape.html
1185    /// [PyArray_Newshape]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_Newshape
1186    fn reshape_with_order<ID: IntoDimension>(
1187        &self,
1188        shape: ID,
1189        order: NPY_ORDER,
1190    ) -> PyResult<Bound<'py, PyArray<T, ID::Dim>>>
1191    where
1192        T: Element;
1193
1194    /// Special case of [`reshape_with_order`][Self::reshape_with_order] which keeps the memory order the same.
1195    #[inline(always)]
1196    fn reshape<ID: IntoDimension>(&self, shape: ID) -> PyResult<Bound<'py, PyArray<T, ID::Dim>>>
1197    where
1198        T: Element,
1199    {
1200        self.reshape_with_order(shape, NPY_ORDER::NPY_ANYORDER)
1201    }
1202
1203    /// Extends or truncates the dimensions of an array.
1204    ///
1205    /// This method works only on [contiguous][PyUntypedArrayMethods::is_contiguous] arrays.
1206    /// Missing elements will be initialized as if calling [`zeros`][PyArray::zeros].
1207    ///
1208    /// See also [`ndarray.resize`][ndarray-resize] and [`PyArray_Resize`][PyArray_Resize].
1209    ///
1210    /// # Safety
1211    ///
1212    /// There should be no outstanding references (shared or exclusive) into the array
1213    /// as this method might re-allocate it and thereby invalidate all pointers into it.
1214    ///
1215    /// # Example
1216    ///
1217    /// ```
1218    /// use numpy::prelude::*;
1219    /// use numpy::PyArray;
1220    /// use pyo3::Python;
1221    ///
1222    /// Python::attach(|py| {
1223    ///     let pyarray = PyArray::<f64, _>::zeros(py, (10, 10), false);
1224    ///     assert_eq!(pyarray.shape(), [10, 10]);
1225    ///
1226    ///     unsafe {
1227    ///         pyarray.resize((100, 100)).unwrap();
1228    ///     }
1229    ///     assert_eq!(pyarray.shape(), [100, 100]);
1230    /// });
1231    /// ```
1232    ///
1233    /// [ndarray-resize]: https://numpy.org/doc/stable/reference/generated/numpy.ndarray.resize.html
1234    /// [PyArray_Resize]: https://numpy.org/doc/stable/reference/c-api/array.html#c.PyArray_Resize
1235    unsafe fn resize<ID: IntoDimension>(&self, newshape: ID) -> PyResult<()>
1236    where
1237        T: Element;
1238
1239    /// Try to convert this array into a [`nalgebra::MatrixView`] using the given shape and strides.
1240    ///
1241    /// # Safety
1242    ///
1243    /// Calling this method invalidates all exclusive references to the internal data, e.g. `ArrayViewMut` or `MatrixSliceMut`.
1244    #[doc(alias = "nalgebra")]
1245    #[cfg(feature = "nalgebra")]
1246    unsafe fn try_as_matrix<R, C, RStride, CStride>(
1247        &self,
1248    ) -> Option<nalgebra::MatrixView<'_, T, R, C, RStride, CStride>>
1249    where
1250        T: nalgebra::Scalar + Element,
1251        D: Dimension,
1252        R: nalgebra::Dim,
1253        C: nalgebra::Dim,
1254        RStride: nalgebra::Dim,
1255        CStride: nalgebra::Dim;
1256
1257    /// Try to convert this array into a [`nalgebra::MatrixViewMut`] using the given shape and strides.
1258    ///
1259    /// # Safety
1260    ///
1261    /// Calling this method invalidates all other references to the internal data, e.g. `ArrayView`, `MatrixSlice`, `ArrayViewMut` or `MatrixSliceMut`.
1262    #[doc(alias = "nalgebra")]
1263    #[cfg(feature = "nalgebra")]
1264    unsafe fn try_as_matrix_mut<R, C, RStride, CStride>(
1265        &self,
1266    ) -> Option<nalgebra::MatrixViewMut<'_, T, R, C, RStride, CStride>>
1267    where
1268        T: nalgebra::Scalar + Element,
1269        D: Dimension,
1270        R: nalgebra::Dim,
1271        C: nalgebra::Dim,
1272        RStride: nalgebra::Dim,
1273        CStride: nalgebra::Dim;
1274}
1275
1276/// Implementation of functionality for [`PyArray0<T>`].
1277#[doc(alias = "PyArray", alias = "PyArray0")]
1278pub trait PyArray0Methods<'py, T>: PyArrayMethods<'py, T, Ix0> {
1279    /// Get the single element of a zero-dimensional array.
1280    ///
1281    /// See [`inner`][crate::inner] for an example.
1282    fn item(&self) -> T
1283    where
1284        T: Element + Copy,
1285    {
1286        unsafe { *self.data() }
1287    }
1288}
1289
1290#[inline(always)]
1291fn get_raw<T, D, Idx>(slf: &Bound<'_, PyArray<T, D>>, index: Idx) -> Option<*mut T>
1292where
1293    T: Element,
1294    D: Dimension,
1295    Idx: NpyIndex<Dim = D>,
1296{
1297    let offset = index.get_checked::<T>(slf.shape(), slf.strides())?;
1298    Some(unsafe { slf.data().offset(offset) })
1299}
1300
1301fn as_view<T, D, S, F>(slf: &Bound<'_, PyArray<T, D>>, from_shape_ptr: F) -> ArrayBase<S, D>
1302where
1303    T: Element,
1304    D: Dimension,
1305    S: RawData,
1306    F: FnOnce(StrideShape<D>, *mut T) -> ArrayBase<S, D>,
1307{
1308    fn inner<D: Dimension>(
1309        shape: &[usize],
1310        strides: &[isize],
1311        itemsize: usize,
1312        mut data_ptr: *mut u8,
1313    ) -> (StrideShape<D>, u32, *mut u8) {
1314        let shape = D::from_dimension(&Dim(shape)).expect(DIMENSIONALITY_MISMATCH_ERR);
1315
1316        assert!(strides.len() <= 32, "{}", MAX_DIMENSIONALITY_ERR);
1317
1318        let mut new_strides = D::zeros(strides.len());
1319        let mut inverted_axes = 0_u32;
1320
1321        for i in 0..strides.len() {
1322            // FIXME(kngwyu): Replace this hacky negative strides support with
1323            // a proper constructor, when it's implemented.
1324            // See https://github.com/rust-ndarray/ndarray/issues/842 for more.
1325            if strides[i] >= 0 {
1326                new_strides[i] = strides[i] as usize / itemsize;
1327            } else {
1328                // Move the pointer to the start position.
1329                data_ptr = unsafe { data_ptr.offset(strides[i] * (shape[i] as isize - 1)) };
1330
1331                new_strides[i] = (-strides[i]) as usize / itemsize;
1332                inverted_axes |= 1 << i;
1333            }
1334        }
1335
1336        (shape.strides(new_strides), inverted_axes, data_ptr)
1337    }
1338
1339    let (shape, mut inverted_axes, data_ptr) = inner(
1340        slf.shape(),
1341        slf.strides(),
1342        mem::size_of::<T>(),
1343        slf.data() as _,
1344    );
1345
1346    let mut array = from_shape_ptr(shape, data_ptr as _);
1347
1348    while inverted_axes != 0 {
1349        let axis = inverted_axes.trailing_zeros() as usize;
1350        inverted_axes &= !(1 << axis);
1351
1352        array.invert_axis(Axis(axis));
1353    }
1354
1355    array
1356}
1357
1358#[cfg(feature = "nalgebra")]
1359fn try_as_matrix_shape_strides<N, D, R, C, RStride, CStride>(
1360    slf: &Bound<'_, PyArray<N, D>>,
1361) -> Option<((R, C), (RStride, CStride))>
1362where
1363    N: nalgebra::Scalar + Element,
1364    D: Dimension,
1365    R: nalgebra::Dim,
1366    C: nalgebra::Dim,
1367    RStride: nalgebra::Dim,
1368    CStride: nalgebra::Dim,
1369{
1370    let ndim = slf.ndim();
1371    let shape = slf.shape();
1372    let strides = slf.strides();
1373
1374    if ndim != 1 && ndim != 2 {
1375        return None;
1376    }
1377
1378    if strides.iter().any(|strides| *strides < 0) {
1379        return None;
1380    }
1381
1382    let rows = shape[0];
1383    let cols = *shape.get(1).unwrap_or(&1);
1384
1385    if R::try_to_usize().map(|expected| rows == expected) == Some(false) {
1386        return None;
1387    }
1388
1389    if C::try_to_usize().map(|expected| cols == expected) == Some(false) {
1390        return None;
1391    }
1392
1393    let row_stride = strides[0] as usize / mem::size_of::<N>();
1394    let col_stride = strides
1395        .get(1)
1396        .map_or(rows, |stride| *stride as usize / mem::size_of::<N>());
1397
1398    if RStride::try_to_usize().map(|expected| row_stride == expected) == Some(false) {
1399        return None;
1400    }
1401
1402    if CStride::try_to_usize().map(|expected| col_stride == expected) == Some(false) {
1403        return None;
1404    }
1405
1406    let shape = (R::from_usize(rows), C::from_usize(cols));
1407
1408    let strides = (
1409        RStride::from_usize(row_stride),
1410        CStride::from_usize(col_stride),
1411    );
1412
1413    Some((shape, strides))
1414}
1415
1416impl<'py, T, D> PyArrayMethods<'py, T, D> for Bound<'py, PyArray<T, D>> {
1417    #[inline(always)]
1418    fn as_untyped(&self) -> &Bound<'py, PyUntypedArray> {
1419        unsafe { self.cast_unchecked() }
1420    }
1421
1422    #[inline(always)]
1423    fn data(&self) -> *mut T {
1424        unsafe { (*self.as_array_ptr()).data.cast() }
1425    }
1426
1427    #[inline(always)]
1428    unsafe fn get(&self, index: impl NpyIndex<Dim = D>) -> Option<&T>
1429    where
1430        T: Element,
1431        D: Dimension,
1432    {
1433        let ptr = get_raw(self, index)?;
1434        Some(&*ptr)
1435    }
1436
1437    #[inline(always)]
1438    unsafe fn get_mut(&self, index: impl NpyIndex<Dim = D>) -> Option<&mut T>
1439    where
1440        T: Element,
1441        D: Dimension,
1442    {
1443        let ptr = get_raw(self, index)?;
1444        Some(&mut *ptr)
1445    }
1446
1447    fn get_owned<Idx>(&self, index: Idx) -> Option<T>
1448    where
1449        T: Element,
1450        D: Dimension,
1451        Idx: NpyIndex<Dim = D>,
1452    {
1453        let element = unsafe { self.get(index) };
1454        element.map(|elem| elem.clone_ref(self.py()))
1455    }
1456
1457    fn to_dyn(&self) -> &Bound<'py, PyArray<T, IxDyn>> {
1458        unsafe { self.cast_unchecked() }
1459    }
1460
1461    fn to_vec(&self) -> Result<Vec<T>, NotContiguousError>
1462    where
1463        T: Element,
1464        D: Dimension,
1465    {
1466        let slice = unsafe { self.as_slice() };
1467        slice.map(|slc| T::vec_from_slice(self.py(), slc))
1468    }
1469
1470    fn try_readonly(&self) -> Result<PyReadonlyArray<'py, T, D>, BorrowError>
1471    where
1472        T: Element,
1473        D: Dimension,
1474    {
1475        PyReadonlyArray::try_new(self.clone())
1476    }
1477
1478    fn try_readwrite(&self) -> Result<PyReadwriteArray<'py, T, D>, BorrowError>
1479    where
1480        T: Element,
1481        D: Dimension,
1482    {
1483        PyReadwriteArray::try_new(self.clone())
1484    }
1485
1486    unsafe fn as_array(&self) -> ArrayView<'_, T, D>
1487    where
1488        T: Element,
1489        D: Dimension,
1490    {
1491        as_view(self, |shape, ptr| ArrayView::from_shape_ptr(shape, ptr))
1492    }
1493
1494    unsafe fn as_array_mut(&self) -> ArrayViewMut<'_, T, D>
1495    where
1496        T: Element,
1497        D: Dimension,
1498    {
1499        as_view(self, |shape, ptr| ArrayViewMut::from_shape_ptr(shape, ptr))
1500    }
1501
1502    fn as_raw_array(&self) -> RawArrayView<T, D>
1503    where
1504        T: Element,
1505        D: Dimension,
1506    {
1507        as_view(self, |shape, ptr| unsafe {
1508            RawArrayView::from_shape_ptr(shape, ptr)
1509        })
1510    }
1511
1512    fn as_raw_array_mut(&self) -> RawArrayViewMut<T, D>
1513    where
1514        T: Element,
1515        D: Dimension,
1516    {
1517        as_view(self, |shape, ptr| unsafe {
1518            RawArrayViewMut::from_shape_ptr(shape, ptr)
1519        })
1520    }
1521
1522    fn to_owned_array(&self) -> Array<T, D>
1523    where
1524        T: Element,
1525        D: Dimension,
1526    {
1527        let view = unsafe { self.as_array() };
1528        T::array_from_view(self.py(), view)
1529    }
1530
1531    fn copy_to<U: Element>(&self, other: &Bound<'py, PyArray<U, D>>) -> PyResult<()>
1532    where
1533        T: Element,
1534    {
1535        let self_ptr = self.as_array_ptr();
1536        let other_ptr = other.as_array_ptr();
1537        let result = unsafe { PY_ARRAY_API.PyArray_CopyInto(self.py(), other_ptr, self_ptr) };
1538        if result != -1 {
1539            Ok(())
1540        } else {
1541            Err(PyErr::fetch(self.py()))
1542        }
1543    }
1544
1545    fn cast_array<U: Element>(&self, is_fortran: bool) -> PyResult<Bound<'py, PyArray<U, D>>>
1546    where
1547        T: Element,
1548    {
1549        let ptr = unsafe {
1550            PY_ARRAY_API.PyArray_CastToType(
1551                self.py(),
1552                self.as_array_ptr(),
1553                U::get_dtype(self.py()).into_dtype_ptr(),
1554                if is_fortran { -1 } else { 0 },
1555            )
1556        };
1557        unsafe { Bound::from_owned_ptr_or_err(self.py(), ptr).map(|ob| ob.cast_into_unchecked()) }
1558    }
1559
1560    fn permute<ID: IntoDimension>(&self, axes: Option<ID>) -> PyResult<Bound<'py, PyArray<T, D>>> {
1561        let mut axes = axes.map(|axes| axes.into_dimension());
1562        let mut axes = axes.as_mut().map(|axes| axes.to_npy_dims());
1563        let axes = axes
1564            .as_mut()
1565            .map_or_else(ptr::null_mut, |axes| axes as *mut npyffi::PyArray_Dims);
1566
1567        let py = self.py();
1568        let ptr = unsafe { PY_ARRAY_API.PyArray_Transpose(py, self.as_array_ptr(), axes) };
1569        unsafe { Bound::from_owned_ptr_or_err(py, ptr).map(|ob| ob.cast_into_unchecked()) }
1570    }
1571
1572    fn reshape_with_order<ID: IntoDimension>(
1573        &self,
1574        shape: ID,
1575        order: NPY_ORDER,
1576    ) -> PyResult<Bound<'py, PyArray<T, ID::Dim>>>
1577    where
1578        T: Element,
1579    {
1580        let mut shape = shape.into_dimension();
1581        let mut shape = shape.to_npy_dims();
1582
1583        let py = self.py();
1584        let ptr = unsafe {
1585            PY_ARRAY_API.PyArray_Newshape(
1586                py,
1587                self.as_array_ptr(),
1588                &mut shape as *mut npyffi::PyArray_Dims,
1589                order,
1590            )
1591        };
1592        unsafe { Bound::from_owned_ptr_or_err(py, ptr).map(|ob| ob.cast_into_unchecked()) }
1593    }
1594
1595    unsafe fn resize<ID: IntoDimension>(&self, newshape: ID) -> PyResult<()>
1596    where
1597        T: Element,
1598    {
1599        let mut newshape = newshape.into_dimension();
1600        let mut newshape = newshape.to_npy_dims();
1601
1602        let py = self.py();
1603        let res = PY_ARRAY_API.PyArray_Resize(
1604            py,
1605            self.as_array_ptr(),
1606            &mut newshape as *mut npyffi::PyArray_Dims,
1607            1,
1608            NPY_ORDER::NPY_ANYORDER,
1609        );
1610
1611        if !res.is_null() {
1612            Ok(())
1613        } else {
1614            Err(PyErr::fetch(py))
1615        }
1616    }
1617
1618    #[cfg(feature = "nalgebra")]
1619    unsafe fn try_as_matrix<R, C, RStride, CStride>(
1620        &self,
1621    ) -> Option<nalgebra::MatrixView<'_, T, R, C, RStride, CStride>>
1622    where
1623        T: nalgebra::Scalar + Element,
1624        D: Dimension,
1625        R: nalgebra::Dim,
1626        C: nalgebra::Dim,
1627        RStride: nalgebra::Dim,
1628        CStride: nalgebra::Dim,
1629    {
1630        let (shape, strides) = try_as_matrix_shape_strides(self)?;
1631
1632        let storage = nalgebra::ViewStorage::from_raw_parts(self.data(), shape, strides);
1633
1634        Some(nalgebra::Matrix::from_data(storage))
1635    }
1636
1637    #[cfg(feature = "nalgebra")]
1638    unsafe fn try_as_matrix_mut<R, C, RStride, CStride>(
1639        &self,
1640    ) -> Option<nalgebra::MatrixViewMut<'_, T, R, C, RStride, CStride>>
1641    where
1642        T: nalgebra::Scalar + Element,
1643        D: Dimension,
1644        R: nalgebra::Dim,
1645        C: nalgebra::Dim,
1646        RStride: nalgebra::Dim,
1647        CStride: nalgebra::Dim,
1648    {
1649        let (shape, strides) = try_as_matrix_shape_strides(self)?;
1650
1651        let storage = nalgebra::ViewStorageMut::from_raw_parts(self.data(), shape, strides);
1652
1653        Some(nalgebra::Matrix::from_data(storage))
1654    }
1655}
1656
1657impl<'py, T> PyArray0Methods<'py, T> for Bound<'py, PyArray0<T>> {}
1658
1659#[cfg(test)]
1660mod tests {
1661    use super::*;
1662
1663    use ndarray::array;
1664    use pyo3::{py_run, types::PyList};
1665
1666    #[test]
1667    fn test_dyn_to_owned_array() {
1668        Python::attach(|py| {
1669            let array = PyArray::from_vec2(py, &[vec![1, 2], vec![3, 4]])
1670                .unwrap()
1671                .to_dyn()
1672                .to_owned_array();
1673
1674            assert_eq!(array, array![[1, 2], [3, 4]].into_dyn());
1675        });
1676    }
1677
1678    #[test]
1679    fn test_hasobject_flag() {
1680        Python::attach(|py| {
1681            let array: Bound<'_, PyArray<Py<PyAny>, _>> =
1682                PyArray1::from_slice(py, &[PyList::empty(py).into()]);
1683
1684            py_run!(py, array, "assert array.dtype.hasobject");
1685        });
1686    }
1687}