Merge pull request #60 from kngwyu/cast

Implement copy and cast methods

Author: kngwyu

src/array.rs

@@ -455,4 +455,95 @@ impl<T: TypeNum> PyArray<T> {
             Self::from_owned_ptr(py, ptr)
         }
     }
+
+    /// Copies self into `other`, performing a data-type conversion if necessary.
+    /// # Example
+    /// ```
+    /// # extern crate pyo3; extern crate numpy; fn main() {
+    /// use numpy::{PyArray, PyArrayModule, IntoPyArray};
+    /// let gil = pyo3::Python::acquire_gil();
+    /// let np = PyArrayModule::import(gil.python()).unwrap();
+    /// let pyarray_f = PyArray::<f64>::arange(gil.python(), &np, 2.0, 5.0, 1.0);
+    /// let mut pyarray_i = PyArray::<i64>::new(gil.python(), &np, &[3]);
+    /// assert!(pyarray_f.copy_to(&np, &mut pyarray_i).is_ok());
+    /// assert_eq!(pyarray_i.as_slice().unwrap(), &[2, 3, 4]);
+    /// # }
+    pub fn copy_to<U: TypeNum>(
+        &self,
+        np: &PyArrayModule,
+        other: &mut PyArray<U>,
+    ) -> Result<(), ArrayCastError> {
+        let self_ptr = self.as_array_ptr();
+        let other_ptr = other.as_array_ptr();
+        let result = unsafe { np.PyArray_CopyInto(other_ptr, self_ptr) };
+        if result == -1 {
+            Err(ArrayCastError::Numpy {
+                from: T::npy_data_type(),
+                to: U::npy_data_type(),
+            })
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Move the data of self into `other`, performing a data-type conversion if necessary.
+    /// # Example
+    /// ```
+    /// # extern crate pyo3; extern crate numpy; fn main() {
+    /// use numpy::{PyArray, PyArrayModule, IntoPyArray};
+    /// let gil = pyo3::Python::acquire_gil();
+    /// let np = PyArrayModule::import(gil.python()).unwrap();
+    /// let pyarray_f = PyArray::<f64>::arange(gil.python(), &np, 2.0, 5.0, 1.0);
+    /// let mut pyarray_i = PyArray::<i64>::new(gil.python(), &np, &[3]);
+    /// assert!(pyarray_f.move_to(&np, &mut pyarray_i).is_ok());
+    /// assert_eq!(pyarray_i.as_slice().unwrap(), &[2, 3, 4]);
+    /// # }
+    pub fn move_to<U: TypeNum>(
+        self,
+        np: &PyArrayModule,
+        other: &mut PyArray<U>,
+    ) -> Result<(), ArrayCastError> {
+        let self_ptr = self.as_array_ptr();
+        let other_ptr = other.as_array_ptr();
+        let result = unsafe { np.PyArray_MoveInto(other_ptr, self_ptr) };
+        if result == -1 {
+            Err(ArrayCastError::Numpy {
+                from: T::npy_data_type(),
+                to: U::npy_data_type(),
+            })
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Cast the `PyArray<T>` to `PyArray<U>`, by allocating a new array.
+    /// # Example
+    /// ```
+    /// # extern crate pyo3; extern crate numpy; fn main() {
+    /// use numpy::{PyArray, PyArrayModule, IntoPyArray};
+    /// let gil = pyo3::Python::acquire_gil();
+    /// let np = PyArrayModule::import(gil.python()).unwrap();
+    /// let pyarray_f = PyArray::<f64>::arange(gil.python(), &np, 2.0, 5.0, 1.0);
+    /// let pyarray_i = pyarray_f.cast::<i32>(gil.python(), &np, false).unwrap();
+    /// assert_eq!(pyarray_i.as_slice().unwrap(), &[2, 3, 4]);
+    /// # }
+    pub fn cast<U: TypeNum>(
+        &self,
+        py: Python,
+        np: &PyArrayModule,
+        is_fortran: bool,
+    ) -> Result<PyArray<U>, ArrayCastError> {
+        let ptr = unsafe {
+            let descr = np.PyArray_DescrFromType(U::typenum_default());
+            np.PyArray_CastToType(self.as_array_ptr(), descr, if is_fortran { -1 } else { 0 })
+        };
+        if ptr.is_null() {
+            Err(ArrayCastError::Numpy {
+                from: T::npy_data_type(),
+                to: U::npy_data_type(),
+            })
+        } else {
+            unsafe { Ok(PyArray::<U>::from_owned_ptr(py, ptr)) }
+        }
+    }
 }

src/error.rs

@@ -28,6 +28,8 @@ pub enum ArrayCastError {
     ToRust { from: NpyDataType, to: NpyDataType },
     /// Error for casting rust's `Vec` into numpy array.
     FromVec,
+    /// Error in numpy -> numpy data conversion
+    Numpy { from: NpyDataType, to: NpyDataType },
 }
 
 impl ArrayCastError {
@@ -46,18 +48,16 @@ impl fmt::Display for ArrayCastError {
                 write!(f, "Cast failed: from={:?}, to={:?}", from, to)
             }
             ArrayCastError::FromVec => write!(f, "Cast failed: FromVec (maybe invalid dimension)"),
+            ArrayCastError::Numpy { from, to } => write!(
+                f,
+                "Cast failed: from=ndarray(dtype={:?}), to=ndarray(dtype={:?})",
+                from, to
+            ),
         }
     }
 }
 
-impl error::Error for ArrayCastError {
-    fn description(&self) -> &str {
-        match self {
-            ArrayCastError::ToRust { .. } => "ArrayCast failed(IntoArray)",
-            ArrayCastError::FromVec => "ArrayCast failed(FromVec)",
-        }
-    }
-}
+impl error::Error for ArrayCastError {}
 
 impl IntoPyErr for ArrayCastError {
     fn into_pyerr(self, msg: &str) -> PyErr {
@@ -67,6 +67,10 @@ impl IntoPyErr for ArrayCastError {
                 from, to, msg
             ),
             ArrayCastError::FromVec => format!("ArrayCastError::FromVec: {}", msg),
+            ArrayCastError::Numpy { from, to } => format!(
+                "ArrayCastError::Numpy: from: {:?}, to: {:?}, msg: {}",
+                from, to, msg
+            ),
         };
         PyErr::new::<exc::TypeError, _>(msg)
     }

tests/array.rs

@@ -181,3 +181,17 @@ macro_rules! small_array_test {
 }
 
 small_array_test!(i8 u8 i16 u16 i32 u32 i64 u64);
+
+#[test]
+fn array_cast() {
+    let gil = pyo3::Python::acquire_gil();
+    let py = gil.python();
+    let np = PyArrayModule::import(py).unwrap();
+    let vec2 = vec![vec![1.0, 2.0, 3.0]; 2];
+    let arr_f64 = PyArray::from_vec2(gil.python(), &np, &vec2).unwrap();
+    let arr_i32: PyArray<i32> = arr_f64.cast(py, &np, false).unwrap();
+    assert_eq!(
+        arr_i32.as_array().unwrap(),
+        array![[1, 2, 3], [1, 2, 3]].into_dyn()
+    );
+}