Revive into_pyarray

Author: kngwyu

example/extensions/src/lib.rs

@@ -3,7 +3,7 @@ extern crate numpy;
 extern crate pyo3;
 
 use ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
-use numpy::{IntoPyResult, PyArrayDyn, ToPyArray};
+use numpy::{IntoPyArray, IntoPyResult, PyArray1, PyArrayDyn, ToPyArray};
 use pyo3::prelude::{pymodinit, PyModule, PyResult, Python};
 
 #[pymodinit]
@@ -28,7 +28,7 @@ fn rust_ext(_py: Python, m: &PyModule) -> PyResult<()> {
     ) -> PyResult<PyArrayDyn<f64>> {
         // you can convert numpy error into PyErr via ?
         let x = x.as_array()?;
-         // you can also specify your error context, via closure
+        // you can also specify your error context, via closure
         let y = y.as_array().into_pyresult_with(|| "y must be f64 array")?;
         Ok(axpy(a, x, y).to_pyarray(py).to_owned(py))
     }
@@ -41,5 +41,15 @@ fn rust_ext(_py: Python, m: &PyModule) -> PyResult<()> {
         Ok(())
     }
 
+    #[pyfn(m, "get_vec")]
+    fn get_vec(py: Python, size: usize) -> PyResult<&PyArray1<f32>> {
+        Ok(vec![0.0; size].into_pyarray(py))
+    }
+    // use numpy::slice_box::SliceBox;
+    // #[pyfn(m, "get_slice")]
+    // fn get_slice(py: Python, size: usize) -> PyResult<SliceBox<f32>> {
+    //     let sbox = numpy::slice_box::SliceBox::new(vec![0.0; size].into_boxed_slice());
+    //     Ok(sbox)
+    // }
     Ok(())
 }

example/setup.py

@@ -10,14 +10,7 @@
 class CmdTest(TestCommand):
     def run(self):
         self.run_command("test_rust")
-        test_files = os.listdir('./tests')
-        ok = 0
-        for f in test_files:
-            _, ext = os.path.splitext(f)
-            if ext == '.py':
-                res = subprocess.call([sys.executable, f], cwd='./tests')
-                ok = ok | res
-        sys.exit(res)
+        subprocess.check_call([sys.executable, 'test_ext.py'], cwd='./tests')
 
 
 setup_requires = ['setuptools-rust>=0.6.0']

example/tests/test_ext.py

@@ -1,5 +1,5 @@
 import numpy as np
-from rust_ext import axpy, mult
+from rust_ext import axpy, mult, get_vec
 import unittest
 
 class TestExt(unittest.TestCase):
@@ -17,6 +17,12 @@ def test_mult(self):
         mult(3.0, x)
         np.testing.assert_array_almost_equal(x, np.array([3.0, 6.0, 9.0]))
 
+    def test_into_pyarray(self):
+        x = get_vec(1000)
+        np.testing.assert_array_almost_equal(x, np.zeros(1000))
+
 
 if __name__ == "__main__":
     unittest.main()
+
+

src/array.rs

@@ -11,6 +11,7 @@ use std::ptr;
 
 use convert::{NpyIndex, ToNpyDims};
 use error::{ErrorKind, IntoPyResult};
+use slice_box::SliceBox;
 use types::{NpyDataType, TypeNum};
 
 /// A safe, static-typed interface for
@@ -329,6 +330,31 @@ impl<T: TypeNum, D: Dimension> PyArray<T, D> {
         Self::from_owned_ptr(py, ptr)
     }
 
+    pub(crate) unsafe fn new_with_data<'py, ID>(
+        py: Python<'py>,
+        dims: ID,
+        strides: *mut npy_intp,
+        slice: &SliceBox<T>,
+    ) -> &'py Self
+    where
+        ID: IntoDimension<Dim = D>,
+    {
+        let dims = dims.into_dimension();
+        let ptr = PY_ARRAY_API.PyArray_New(
+            PY_ARRAY_API.get_type_object(npyffi::ArrayType::PyArray_Type),
+            dims.ndim_cint(),
+            dims.as_dims_ptr(),
+            T::typenum_default(),
+            strides,                // strides
+            slice.data(),           // data
+            0,                      // itemsize
+            0,                      // flag
+            ::std::ptr::null_mut(), //obj
+        );
+        PY_ARRAY_API.PyArray_SetBaseObject(ptr as *mut npyffi::PyArrayObject, slice.as_ptr());
+        Self::from_owned_ptr(py, ptr)
+    }
+
     /// Construct a new nd-dimensional array filled with 0.
     ///
     /// If `is_fortran` is true, then

src/convert.rs

@@ -2,13 +2,54 @@
 
 use ndarray::{ArrayBase, Data, Dimension, IntoDimension, Ix1};
 use pyo3::Python;
+use slice_box::SliceBox;
 
 use std::mem;
 use std::os::raw::c_int;
+use std::ptr;
 
 use super::*;
 
-/// Covversion trait from rust types to `PyArray`.
+/// Covnersion trait from some rust types to `PyArray`.
+///
+/// This trait takes `self`, which means **it holds a pointer to Rust heap, until `resize` or other
+/// destructive method is called**.
+/// # Example
+/// ```
+/// # extern crate pyo3; extern crate numpy; fn main() {
+/// use numpy::{PyArray, ToPyArray};
+/// let gil = pyo3::Python::acquire_gil();
+/// let py_array = vec![1, 2, 3].to_pyarray(gil.python());
+/// assert_eq!(py_array.as_slice().unwrap(), &[1, 2, 3]);
+/// # }
+/// ```
+pub trait IntoPyArray {
+    type Item: TypeNum;
+    type Dim: Dimension;
+    fn into_pyarray<'py>(self, Python<'py>) -> &'py PyArray<Self::Item, Self::Dim>;
+}
+
+impl<T: TypeNum> IntoPyArray for Box<[T]> {
+    type Item = T;
+    type Dim = Ix1;
+    fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
+        let len = self.len();
+        unsafe {
+            let slice = SliceBox::new(self);
+            PyArray::new_with_data(py, [len], ptr::null_mut(), slice)
+        }
+    }
+}
+
+impl<T: TypeNum> IntoPyArray for Vec<T> {
+    type Item = T;
+    type Dim = Ix1;
+    fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
+        self.into_boxed_slice().into_pyarray(py)
+    }
+}
+
+/// Conversion trait from rust types to `PyArray`.
 ///
 /// This trait takes `&self`, which means **it alocates in Python heap and then copies
 /// elements there**.

src/lib.rs

@@ -44,13 +44,14 @@ pub mod array;
 pub mod convert;
 pub mod error;
 pub mod npyffi;
+mod slice_box;
 pub mod types;
 
 pub use array::{
     get_array_module, PyArray, PyArray1, PyArray2, PyArray3, PyArray4, PyArray5, PyArray6,
     PyArrayDyn,
 };
-pub use convert::{NpyIndex, ToNpyDims, ToPyArray};
+pub use convert::{IntoPyArray, NpyIndex, ToNpyDims, ToPyArray};
 pub use error::{ArrayFormat, ErrorKind, IntoPyErr, IntoPyResult};
 pub use ndarray::{Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
 pub use npyffi::{PY_ARRAY_API, PY_UFUNC_API};

src/slice_box.rs

@@ -0,0 +1,79 @@
+use crate::types::TypeNum;
+use pyo3::{self, ffi, typeob, PyObjectAlloc, Python, ToPyPointer};
+use std::os::raw::c_void;
+
+#[repr(C)]
+pub(crate) struct SliceBox<T> {
+    ob_base: ffi::PyObject,
+    inner: *mut [T],
+}
+
+impl<T> SliceBox<T> {
+    pub(crate) unsafe fn new<'a>(box_: Box<[T]>) -> &'a Self {
+        <Self as typeob::PyTypeObject>::init_type();
+        let type_ob = <Self as typeob::PyTypeInfo>::type_object() as *mut _;
+        let base = ffi::_PyObject_New(type_ob);
+        *base = ffi::PyObject_HEAD_INIT;
+        (*base).ob_type = <Self as typeob::PyTypeInfo>::type_object() as *mut _;
+        let self_ = base as *mut SliceBox<T>;
+        (*self_).inner = Box::into_raw(box_);
+        &*self_
+    }
+    pub(crate) fn data(&self) -> *mut c_void {
+        self.inner as *mut c_void
+    }
+}
+
+impl<T> typeob::PyTypeInfo for SliceBox<T> {
+    type Type = ();
+    type BaseType = pyo3::PyObjectRef;
+    const NAME: &'static str = "SliceBox";
+    const DESCRIPTION: &'static str = "Memory store for PyArray made by IntoPyArray.";
+    const FLAGS: usize = 0;
+    const SIZE: usize = { Self::OFFSET as usize + std::mem::size_of::<Self>() + 0 + 0 };
+    const OFFSET: isize = 0;
+    #[inline]
+    unsafe fn type_object() -> &'static mut ::pyo3::ffi::PyTypeObject {
+        static mut TYPE_OBJECT: ::pyo3::ffi::PyTypeObject = ::pyo3::ffi::PyTypeObject_INIT;
+        &mut TYPE_OBJECT
+    }
+}
+
+impl<T: TypeNum> typeob::PyTypeObject for SliceBox<T> {
+    #[inline(always)]
+    fn init_type() {
+        static START: std::sync::Once = std::sync::ONCE_INIT;
+        START.call_once(|| {
+            let ty = unsafe { <Self as typeob::PyTypeInfo>::type_object() };
+            if (ty.tp_flags & ffi::Py_TPFLAGS_READY) == 0 {
+                let gil = Python::acquire_gil();
+                let py = gil.python();
+                let mod_name = format!("rust_numpy.{:?}", T::npy_data_type());
+                typeob::initialize_type::<Self>(py, Some(&mod_name))
+                    .map_err(|e| e.print(py))
+                    .expect("Failed to initialize SliceBox");
+            }
+        });
+    }
+}
+
+impl<T> ToPyPointer for SliceBox<T> {
+    #[inline]
+    fn as_ptr(&self) -> *mut ffi::PyObject {
+        &self.ob_base as *const _ as *mut _
+    }
+}
+
+impl<T> PyObjectAlloc<SliceBox<T>> for SliceBox<T> {
+    /// Calls the rust destructor for the object.
+    unsafe fn drop(py: Python, obj: *mut ffi::PyObject) {
+        let data = (*(obj as *mut SliceBox<T>)).inner;
+        let box_ = Box::from_raw(data);
+        drop(box_);
+        <Self as typeob::PyTypeInfo>::BaseType::drop(py, obj);
+    }
+    unsafe fn dealloc(py: Python, obj: *mut ffi::PyObject) {
+        Self::drop(py, obj);
+        ffi::PyObject_Free(obj as *mut c_void);
+    }
+}

tests/array.rs

@@ -75,7 +75,7 @@ fn as_array() {
 }
 
 #[test]
-fn into_pyarray_vec() {
+fn to_pyarray_vec() {
     let gil = pyo3::Python::acquire_gil();
 
     let a = vec![1, 2, 3];
@@ -86,7 +86,7 @@ fn into_pyarray_vec() {
 }
 
 #[test]
-fn into_pyarray_array() {
+fn to_pyarray_array() {
     let gil = pyo3::Python::acquire_gil();
 
     let a = Array3::<f64>::zeros((3, 4, 2));
@@ -206,3 +206,11 @@ fn array_cast() {
     let arr_i32: &PyArray2<i32> = arr_f64.cast(false).unwrap();
     assert_eq!(arr_i32.as_array().unwrap(), array![[1, 2, 3], [1, 2, 3]]);
 }
+
+#[test]
+fn into_pyarray_vec() {
+    let gil = pyo3::Python::acquire_gil();
+    let a = vec![1, 2, 3];
+    let arr = a.into_pyarray(gil.python());
+    assert_eq!(arr.as_slice().unwrap(), &[1, 2, 3])
+}