array now returns a reference in the original vector instead of a copy

Author: dirkvdb

infra-rs

@@ -93,7 +93,8 @@ class raster:
     @property
     def array(self) -> np.ndarray:
         """
-        Returns the raster as a numpy array. Data is copied so changes to the array will not affect the raster.
+        Returns the raster as a numpy array. The data in the array is a view on the raster data so changes to the array will affect the raster.
+        This also means that the array cannot be changed in size, only in content.
         """
         ...
     @property

python/geodynamix/geodynamix.pyi

@@ -1,5 +1,9 @@
 use numpy::PyArrayDescr;
+use pyo3::exceptions::PyBufferError;
+use pyo3::ffi;
 use pyo3::prelude::*;
+use std::ffi::c_void;
+use std::ffi::CString;
 use std::ops::Add;
 use std::ops::Div;
 use std::ops::Mul;
@@ -71,8 +75,9 @@ impl Raster {
     }
 
     #[getter]
-    pub fn array<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        utils::raster_array(py, &self.raster)
+    pub fn array<'py>(slf: Bound<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        //utils::raster_array(py, &slf.borrow().raster)
+        utils::raster_buffer_array(py, slf)
     }
 
     #[getter]
@@ -109,6 +114,26 @@ impl Raster {
     pub fn __truediv__(&self, rhs_object: Bound<'_, PyAny>) -> PyResult<Raster> {
         impl_raster_op!(div, self.raster, rhs_object)
     }
+    /// # Safety
+    /// This function is unsafe because it exposes a raw pointer to the Python buffer protocol.
+    pub unsafe fn __getbuffer__(
+        slf: Bound<'_, Self>,
+        view: *mut ffi::Py_buffer,
+        flags: std::ffi::c_int,
+    ) -> PyResult<()> {
+        fill_view_from_data(
+            view,
+            flags,
+            slf.borrow().raster.raw_data_u8_slice(),
+            slf.into_any(),
+        )
+    }
+
+    /// # Safety
+    /// This function is unsafe because it exposes a raw pointer to the Python buffer protocol.
+    pub unsafe fn __releasebuffer__(&self, _buffer: *mut ffi::Py_buffer) {
+        // No need to release, the raster still owns the data
+    }
 }
 
 impl From<PythonDenseArray> for Raster {
@@ -144,3 +169,50 @@ pub fn raster_equal(
 
     Ok(array1 == array2)
 }
+
+/// # Safety
+///
+/// `view` must be a valid pointer to `ffi::Py_buffer`, or null
+/// `data` must outlive the Python lifetime of `owner` (i.e. data must be owned by owner, or data
+/// must be static data)
+unsafe fn fill_view_from_data(
+    view: *mut ffi::Py_buffer,
+    flags: std::ffi::c_int,
+    data: &[u8],
+    owner: Bound<'_, PyAny>,
+) -> PyResult<()> {
+    if view.is_null() {
+        return Err(PyBufferError::new_err("View is null"));
+    }
+
+    (*view).obj = owner.into_ptr();
+    (*view).buf = data.as_ptr() as *mut c_void;
+    (*view).len = data.len() as isize;
+    (*view).readonly = 0;
+    (*view).itemsize = 1;
+
+    (*view).format = if (flags & ffi::PyBUF_FORMAT) == ffi::PyBUF_FORMAT {
+        let msg = CString::new("B").unwrap();
+        msg.into_raw()
+    } else {
+        std::ptr::null_mut()
+    };
+
+    (*view).ndim = 1;
+    (*view).shape = if (flags & ffi::PyBUF_ND) == ffi::PyBUF_ND {
+        &mut (*view).len
+    } else {
+        std::ptr::null_mut()
+    };
+
+    (*view).strides = if (flags & ffi::PyBUF_STRIDES) == ffi::PyBUF_STRIDES {
+        &mut (*view).itemsize
+    } else {
+        std::ptr::null_mut()
+    };
+
+    (*view).suboffsets = std::ptr::null_mut();
+    (*view).internal = std::ptr::null_mut();
+
+    Ok(())
+}

src/raster.rs

@@ -3,7 +3,7 @@ use pyo3::{exceptions::PyValueError, prelude::*, types::IntoPyDict};
 
 use geo::{raster::algo, Array, ArrayDataType, ArrayMetadata, ArrayNum, DenseArray};
 
-use crate::{PythonDenseArray, RasterMetadata};
+use crate::{PythonDenseArray, Raster, RasterMetadata};
 
 fn convert_array<'py, T: ArrayNum + numpy::Element>(
     py: Python<'py>,
@@ -47,6 +47,27 @@ pub fn raster_mask<'py>(py: Python<'py>, raster: &PythonDenseArray) -> PyResult<
     }
 }
 
+pub fn raster_buffer_array<'py>(
+    py: Python<'py>,
+    raster: Bound<'py, Raster>,
+) -> PyResult<Bound<'py, PyAny>> {
+    let np = PyModule::import(py, "numpy")?;
+
+    let shape = [
+        raster.borrow().raster.rows().count(),
+        raster.borrow().raster.columns().count(),
+    ];
+    let dtype =
+        array_type_to_numpy_dtype(py, raster.borrow().raster.data_type()).into_pyobject(py)?;
+    let kwargs = vec![
+        ("dtype", dtype.into_pyobject(py)?.into_any()),
+        ("buffer", raster.into_pyobject(py)?.into_any()),
+    ];
+
+    np.getattr("ndarray")?
+        .call((shape.into_pyobject(py)?,), Some(&kwargs.into_py_dict(py)?))
+}
+
 pub fn raster_masked_array<'py>(
     py: Python<'py>,
     raster: &PythonDenseArray,

src/utils.rs

@@ -182,31 +182,31 @@ def test_write_map(self):
         os.remove("writtenraster.asc")
         self.assertTrue(np.allclose(raster, written_raster.array, equal_nan=True))
 
-    # def test_modify_raster_using_numpy(self):
-    #     # modify the raster data using the ndarry accessor
-    #     # verify that the internal raster data has changed by writing it to disk
-    #     # and comparing the result
+    def test_modify_raster_using_numpy(self):
+        # modify the raster data using the ndarry accessor
+        # verify that the internal raster data has changed by writing it to disk
+        # and comparing the result
 
-    #     create_test_file("raster.asc", test_raster)
-    #     raster = gdx.read("raster.asc")
+        create_test_file("raster.asc", test_raster)
+        raster = gdx.read("raster.asc")
 
-    #     raster.array.fill(44)
+        raster.array.fill(44)
 
-    #     gdx.write(raster, "writtenraster.asc")
-    #     written_raster = gdx.read("writtenraster.asc")
-    #     os.remove("writtenraster.asc")
+        gdx.write(raster, "writtenraster.asc")
+        written_raster = gdx.read("writtenraster.asc")
+        os.remove("writtenraster.asc")
 
-    #     expected = np.array(
-    #         [
-    #             [44, 44, 44, 44, 44],
-    #             [44, 44, 44, 44, 44],
-    #             [44, 44, 44, 44, 44],
-    #             [44, 44, 44, 44, 44],
-    #         ],
-    #         dtype="f",
-    #     )
+        expected = np.array(
+            [
+                [44, 44, 44, 44, 44],
+                [44, 44, 44, 44, 44],
+                [44, 44, 44, 44, 44],
+                [44, 44, 44, 44, 44],
+            ],
+            dtype="f",
+        )
 
-    #     self.assertTrue(np.allclose(expected, written_raster.array, equal_nan=True))
+        self.assertTrue(np.allclose(expected, written_raster.array, equal_nan=True))
 
     def test_replace_value(self):
         expected = np.array(