Add a test case showing how to extract an array from a dictionary.

Author: adamreichold

examples/simple-extension/src/lib.rs

@@ -1,6 +1,10 @@
 use numpy::ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
-use numpy::{Complex64, IntoPyArray, PyArrayDyn, PyReadonlyArrayDyn};
-use pyo3::{pymodule, types::PyModule, PyResult, Python};
+use numpy::{Complex64, IntoPyArray, PyArray1, PyArrayDyn, PyReadonlyArrayDyn};
+use pyo3::{
+    pymodule,
+    types::{PyDict, PyModule},
+    PyResult, Python,
+};
 
 #[pymodule]
 fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
@@ -52,5 +56,17 @@ fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
         conj(x.as_array()).into_pyarray(py)
     }
 
+    #[pyfn(m)]
+    #[pyo3(name = "extract")]
+    fn extract(d: &PyDict) -> f64 {
+        let x = d
+            .get_item("x")
+            .unwrap()
+            .downcast::<PyArray1<f64>>()
+            .unwrap();
+
+        x.readonly().as_array().sum()
+    }
+
     Ok(())
 }

examples/simple-extension/tests/test_ext.py

@@ -1,5 +1,5 @@
 import numpy as np
-from rust_ext import axpy, conj, mult
+from rust_ext import axpy, conj, mult, extract
 
 
 def test_axpy():
@@ -22,3 +22,9 @@ def test_mult():
 def test_conj():
     x = np.array([1.0 + 2j, 2.0 + 3j, 3.0 + 4j])
     np.testing.assert_array_almost_equal(conj(x), np.conj(x))
+
+
+def test_extract():
+    x = np.arange(5)
+    d = { "x": x }
+    np.testing.assert_almost_equal(extract(d), 10.0)