use crate::npyffi::NPY_ARRAY_WRITEABLE;
use crate::{Element, NotContiguousError, NpyIndex, PyArray};
use ndarray::{ArrayView, Dimension, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
use pyo3::{prelude::*, types::PyAny, AsPyPointer};
pub struct PyReadonlyArray<'py, T, D> {
array: &'py PyArray<T, D>,
was_writeable: bool,
}
impl<'py, T: Element, D: Dimension> PyReadonlyArray<'py, T, D> {
pub fn as_slice(&self) -> Result<&[T], NotContiguousError> {
unsafe { self.array.as_slice() }
}
pub fn as_array(&self) -> ArrayView<'_, T, D> {
unsafe { self.array.as_array() }
}
#[inline(always)]
pub fn get(&self, index: impl NpyIndex<Dim = D>) -> Option<&T> {
unsafe { self.array.get(index) }
}
}
pub type PyReadonlyArray1<'py, T> = PyReadonlyArray<'py, T, Ix1>;
pub type PyReadonlyArray2<'py, T> = PyReadonlyArray<'py, T, Ix2>;
pub type PyReadonlyArray3<'py, T> = PyReadonlyArray<'py, T, Ix3>;
pub type PyReadonlyArray4<'py, T> = PyReadonlyArray<'py, T, Ix4>;
pub type PyReadonlyArray5<'py, T> = PyReadonlyArray<'py, T, Ix5>;
pub type PyReadonlyArray6<'py, T> = PyReadonlyArray<'py, T, Ix6>;
pub type PyReadonlyArrayDyn<'py, T> = PyReadonlyArray<'py, T, IxDyn>;
impl<'py, T: Element, D: Dimension> FromPyObject<'py> for PyReadonlyArray<'py, T, D> {
fn extract(obj: &'py PyAny) -> PyResult<Self> {
let array: &PyArray<T, D> = obj.extract()?;
Ok(PyReadonlyArray::from(array))
}
}
impl<'py, T, D> IntoPy<PyObject> for PyReadonlyArray<'py, T, D> {
fn into_py(self, py: Python<'_>) -> PyObject {
let PyReadonlyArray { array, .. } = self;
unsafe { PyObject::from_borrowed_ptr(py, array.as_ptr()) }
}
}
impl<'py, T, D> From<&'py PyArray<T, D>> for PyReadonlyArray<'py, T, D> {
fn from(array: &'py PyArray<T, D>) -> PyReadonlyArray<'py, T, D> {
let flag = array.get_flag();
let writeable = flag & NPY_ARRAY_WRITEABLE != 0;
if writeable {
unsafe {
(*array.as_array_ptr()).flags &= !NPY_ARRAY_WRITEABLE;
}
}
Self {
array,
was_writeable: writeable,
}
}
}
impl<'py, T, D> Drop for PyReadonlyArray<'py, T, D> {
fn drop(&mut self) {
if self.was_writeable {
unsafe {
(*self.array.as_array_ptr()).flags |= NPY_ARRAY_WRITEABLE;
}
}
}
}
impl<'py, T, D> AsRef<PyArray<T, D>> for PyReadonlyArray<'py, T, D> {
fn as_ref(&self) -> &PyArray<T, D> {
self.array
}
}
impl<'py, T, D> std::ops::Deref for PyReadonlyArray<'py, T, D> {
type Target = PyArray<T, D>;
fn deref(&self) -> &PyArray<T, D> {
self.array
}
}