#numpy #python #pyo3 #ndarray

pyo3-arraylike

rust-numpy的便利扩展

2 个版本

0.1.1 2024年4月13日
0.1.0 2023年7月11日

#109 in FFI

MIT 许可证

14KB
250

PyO3-ArrayLike

此crate提供了一个名为 PyArrayLike<T,D> 的单结构体,可以用于从任何可以被视为类型为 T 和维度 D 的数组类型的Python对象中合理地提取数组。

工作原理

PyArrayLike<T,D>基本上是一个枚举,由来自 rust-numpy 项目的 PyReadonlyArray<T,D>(即对Python堆上实际NumPy数组的只读引用)或来自 ndarray 项目的所有者 Array<T,D> 组成。当调用 pyobject.extract::<PyArrayLike<T,D>>() 时,它将立即通过内部调用 pyobject.downcast::<PyReadonlyArray<T,D>>() 来检查 pyobject 是否已经是正确类型和维度的NumPy数组。只有当这失败时,它才会尝试通过遍历其元素并将它们递归地转换为正确类型来从 pyobject 中构建一个所有者数组。

示例

考虑以下函数。它接受任何看起来像由整数(0到4294967295之间)组成的二维数组的对象。它计算数组的行和,并返回一个一维NumPy数组。

#[pyfunction]
fn sum_of_rows<'py>(py: Python<'py>, ar: PyArrayLike2<'py, u32>) -> Bound<'py, PyArray1<u32>> {
    ar.view().sum_axis(Axis(0)).into_pyarray_bound(py)
}

您可以从Python以以下方式调用此函数。前两种方法将成功,后两种将引发错误。

def call1():
    """Succeeds by passing a reference of the input array to the rust component."""
    return sum_of_rows(np.array([[1,2,3],[4,5,6]], dtype='uint32'))

def call2():
    """Succeeds but allocates extra memory for rebuilding the input array inside the rust component."""
    return sum_of_rows([[1,2,3], iter([4,5,6]), np.array([7,8,9], dtype='int64')])

def call3():
    """Raises an error since the input array is expected to be two dimensional."""
    return sum_of_rows([1,2,3])

def call4():
    """Raises an error since the input array contains a value which cannot be safely casted to u32."""
    return sum_of_rows([[2**32,0,0], [0,0,0]])

依赖项

~6-11MB
~132K SLoC