tutorial-rust-for-python/src/python_bindings.rs at main · wenig/tutorial-rust-for-python

Repository navigation

use ndarray::{ArrayView1, ArrayView2};

use crate::knn;

use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1, PyReadonlyArray2};

use pyo3::prelude::*;

#[pyfunction]

fn knn_algorithm<'py>(

py: Python<'py>,

x_py: PyReadonlyArray2<'py, f64>,

y_py: PyReadonlyArray1<'py, i64>,

k: usize

) -> PyResult<&'py PyArray1<i64>> {

let x: ArrayView2<f64> = x_py.as_array();

let y: ArrayView1<i64> = y_py.as_array();

let predicted = knn(x, y, k);

Ok(predicted.into_pyarray(py))

}

#[pymodule]

fn knn_rs(_py: Python, m: &PyModule) -> PyResult<()> {

m.add_function(wrap_pyfunction!(knn_algorithm, m)?)?;

Ok(())

}