Error handling, other improvements

argmin-rs · Mar 8, 2024 · 98557bc · 98557bc
1 parent fb810f6
commit 98557bc
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 61 deletions.
diff --git a/python/finitediff-py/Cargo.toml b/python/finitediff-py/Cargo.toml
@@ -21,4 +21,5 @@ crate-type = ["cdylib"]
 [dependencies]
 finitediff_rust = { package = "finitediff", version = "0.1.4", path = "../../crates/finitediff", features = ["ndarray"] }
 numpy = "0.20.0"
-pyo3 = { version = "0.20.0", features = ["extension-module"] }
+pyo3 = { version = "0.20.0", features = ["extension-module", "anyhow"] }
+anyhow = "1.0"
diff --git a/python/finitediff-py/src/lib.rs b/python/finitediff-py/src/lib.rs
@@ -1,12 +1,29 @@
+use anyhow::Error;
 use finitediff_rust::ndarr;
-use numpy::ndarray::Array1;
-use numpy::{IntoPyArray, PyArray1, PyArray2};
+use numpy::{ndarray::Array1, IntoPyArray, PyArray1, PyArray2};
 use pyo3::{
     exceptions::PyTypeError,
     prelude::*,
     types::{PyCFunction, PyDict, PyTuple},
 };
 
+fn process_args(args: &PyTuple) -> PyResult<Array1<f64>> {
+    Ok(args
+        .get_item(0)
+        .map_err(|_| PyErr::new::<PyTypeError, _>("Insufficient number of arguments"))?
+        .downcast::<PyArray1<f64>>()?
+        .to_owned_array())
+}
+
+macro_rules! not_callable {
+    ($py:ident, $f:ident) => {
+        Err(PyErr::new::<PyTypeError, _>(format!(
+            "object {} not callable",
+            $f.as_ref($py).get_type()
+        )))
+    };
+}
+
 /// Forward diff
 #[pyfunction]
 fn forward_diff<'py>(py: Python<'py>, f: Py<PyAny>) -> PyResult<&'py PyCFunction> {
@@ -17,24 +34,16 @@ fn forward_diff<'py>(py: Python<'py>, f: Py<PyAny>) -> PyResult<&'py PyCFunction
             None,
             move |args: &PyTuple, _kwargs: Option<&PyDict>| -> PyResult<Py<PyArray1<f64>>> {
                 Python::with_gil(|py| {
-                    let out = (ndarr::forward_diff(|x: &Array1<f64>| -> f64 {
+                    let out = (ndarr::forward_diff(&|x: &Array1<f64>| -> Result<f64, Error> {
                         let x = PyArray1::from_array(py, x);
-                        f.call(py, (x,), None).unwrap().extract(py).unwrap()
-                    }))(
-                        &args
-                            .get_item(0)?
-                            .downcast::<PyArray1<f64>>()?
-                            .to_owned_array(),
-                    );
+                        Ok(f.call(py, (x,), None)?.extract(py)?)
+                    }))(&process_args(args)?)?;
                     Ok(out.into_pyarray(py).into())
                 })
             },
         )
     } else {
-        Err(PyErr::new::<PyTypeError, _>(format!(
-            "object {} not callable",
-            f.as_ref(py).get_type()
-        )))
+        not_callable!(py, f)
     }
 }
 
@@ -48,24 +57,16 @@ fn central_diff<'py>(py: Python<'py>, f: Py<PyAny>) -> PyResult<&'py PyCFunction
             None,
             move |args: &PyTuple, _kwargs: Option<&PyDict>| -> PyResult<Py<PyArray1<f64>>> {
                 Python::with_gil(|py| {
-                    let out = (ndarr::central_diff(|x: &Array1<f64>| -> f64 {
+                    let out = (ndarr::central_diff(&|x: &Array1<f64>| -> Result<f64, Error> {
                         let x = PyArray1::from_array(py, x);
-                        f.call(py, (x,), None).unwrap().extract(py).unwrap()
-                    }))(
-                        &args
-                            .get_item(0)?
-                            .downcast::<PyArray1<f64>>()?
-                            .to_owned_array(),
-                    );
+                        Ok(f.call(py, (x,), None)?.extract(py)?)
+                    }))(&process_args(args)?)?;
                     Ok(out.into_pyarray(py).into())
                 })
             },
         )
     } else {
-        Err(PyErr::new::<PyTypeError, _>(format!(
-            "object {} not callable",
-            f.as_ref(py).get_type()
-        )))
+        not_callable!(py, f)
     }
 }
 
@@ -79,28 +80,20 @@ fn forward_jacobian<'py>(py: Python<'py>, f: Py<PyAny>) -> PyResult<&'py PyCFunc
             None,
             move |args: &PyTuple, _kwargs: Option<&PyDict>| -> PyResult<Py<PyArray2<f64>>> {
                 Python::with_gil(|py| {
-                    let out = (ndarr::forward_jacobian(|x: &Array1<f64>| -> Array1<f64> {
-                        let x = PyArray1::from_array(py, x);
-                        f.call(py, (x,), None)
-                            .unwrap()
-                            .extract::<&PyArray1<f64>>(py)
-                            .unwrap()
-                            .to_owned_array()
-                    }))(
-                        &args
-                            .get_item(0)?
-                            .downcast::<PyArray1<f64>>()?
-                            .to_owned_array(),
-                    );
+                    let out = (ndarr::forward_jacobian(
+                        &|x: &Array1<f64>| -> Result<Array1<f64>, Error> {
+                            let x = PyArray1::from_array(py, x);
+                            Ok(f.call(py, (x,), None)?
+                                .extract::<&PyArray1<f64>>(py)?
+                                .to_owned_array())
+                        },
+                    ))(&process_args(args)?)?;
                     Ok(out.into_pyarray(py).into())
                 })
             },
         )
     } else {
-        Err(PyErr::new::<PyTypeError, _>(format!(
-            "object {} not callable",
-            f.as_ref(py).get_type()
-        )))
+        not_callable!(py, f)
     }
 }
 
@@ -114,28 +107,20 @@ fn central_jacobian<'py>(py: Python<'py>, f: Py<PyAny>) -> PyResult<&'py PyCFunc
             None,
             move |args: &PyTuple, _kwargs: Option<&PyDict>| -> PyResult<Py<PyArray2<f64>>> {
                 Python::with_gil(|py| {
-                    let out = (ndarr::central_jacobian(|x: &Array1<f64>| -> Array1<f64> {
-                        let x = PyArray1::from_array(py, x);
-                        f.call(py, (x,), None)
-                            .unwrap()
-                            .extract::<&PyArray1<f64>>(py)
-                            .unwrap()
-                            .to_owned_array()
-                    }))(
-                        &args
-                            .get_item(0)?
-                            .downcast::<PyArray1<f64>>()?
-                            .to_owned_array(),
-                    );
+                    let out = (ndarr::central_jacobian(
+                        &|x: &Array1<f64>| -> Result<Array1<f64>, Error> {
+                            let x = PyArray1::from_array(py, x);
+                            Ok(f.call(py, (x,), None)?
+                                .extract::<&PyArray1<f64>>(py)?
+                                .to_owned_array())
+                        },
+                    ))(&process_args(args)?)?;
                     Ok(out.into_pyarray(py).into())
                 })
             },
         )
     } else {
-        Err(PyErr::new::<PyTypeError, _>(format!(
-            "object {} not callable",
-            f.as_ref(py).get_type()
-        )))
+        not_callable!(py, f)
     }
 }
 

diff --git a/python/finitediff-py/test.py b/python/finitediff-py/test.py
@@ -56,6 +56,7 @@ def op(x):
 x = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
 print(j(x))
 
+
 # class NotCallable:
 #     pass
-Original file line number
+Diff line change
@@ Expand Up / @@ -56,6 +56,7 @@ def op(x): @@
     x = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
     print(j(x))
     # class NotCallable:
     #     pass
@@ Expand Down @@