PyO3 / rust-numpy

PyO3-based Rust bindings of the NumPy C-API
BSD 2-Clause "Simplified" License
1.11k stars 106 forks source link

Feature request: implicit conversion #382

Closed 124C41p closed 1 year ago

124C41p commented 1 year ago

When writing a Python function processing numpy arrays, it is often a common practice to accept any "array like" input, and convert it into a numpy array if necessary. For example, the following function processes a float64-array (and it does not allocate memory if such an array is given), but also works for integer lists, or something like that:

def my_mean(ar):
    ar = np.asarray(ar, dtype=np.float64)
    return ar.mean()

I would like to do the same thing in rust with PyO3. The following function should accept anything which can be reasonably converted into a one dimensional f64-array, but should not allocate memory if the "correct" type is plugged in:

#[pyfunction]
fn my_mean(ar: ArrayLike1<f64>) -> PyResult<f64> {
    let res = ar
        .as_array()
        .mean()
        .ok_or(PyValueError::new_err("Empty array"))?;
    Ok(res)
}

Technically, it could be achieved like that:

enum ArrayLike1<'py, T: Element> {
    Reference(PyReadonlyArray1<'py, T>),
    Converted(Array1<T>),
}

impl<'py, T> FromPyObject<'py> for ArrayLike1<'py, T>
where
    T: Element,
    Vec<T>: FromPyObject<'py>,
{
    fn extract(ob: &'py PyAny) -> PyResult<Self> {
        if let Ok(res) = ob.extract::<PyReadonlyArray1<T>>() {
            return Ok(ArrayLike1::Reference(res));
        };

        let res = ob.extract::<Vec<T>>()?;
        Ok(ArrayLike1::Converted(res.into_iter().collect()))
    }
}

impl<'py, T: Element> ArrayLike1<'py, T> {
    pub fn as_array(&self) -> ArrayView1<T> {
        match self {
            ArrayLike1::Reference(x) => x.as_array(),
            ArrayLike1::Converted(x) => x.view(),
        }
    }
}
adamreichold commented 1 year ago

Since you already posted code, why not turn this into a pull request? (The approach itself is reasonable as far I can see.)

124C41p commented 1 year ago

Ok, I'll try. It might take some time however to figure out the general case (say ArrayLike<T, D>) as I am not an expert in any of the involved libraries.

adamreichold commented 1 year ago

Take your time. If you get stuck, you can open a draft pull request with what you have even if it does not build. I will try to help out with open issues then.