问题描述
我正在尝试用 rust 编写一个函数,我可以从 python 中调用它,该函数接受一个字典列表(想想类似 Pandas 数据框的数据)并从 rust 访问这些键值。我怎样才能做到这一点?我正在使用 pyo3
。我是否需要定义一个与 python dict 输入的键、值对匹配的结构?
作为示例函数,我试图传入一个字典列表,并将与键 key
对应的值相加为总数。我的 python dicts 列表中的每个字典都有对应于 int 的键 key
。
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use pyo3::types::PyDict;
#[pyfunction]
fn sum_list_dicts(a: Vec<PyDict>,key: String) -> PyResult<i32> {
let mut tot = 0_i32;
for d in a.iter() {
tot += d[key];
}
Ok(tot)
}
#[pymodule]
fn rustpy(_py: Python,m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(sum_list_dicts,m)?)?;
Ok(())
}
解决方法
所以这真的取决于你真正想要做什么。如果您不想弄乱实际的 Py
项,您可以简单地这样做:
#[pyfunction]
fn sum_list_dicts(a: Vec<HashMap<String,i32>>,key: String) -> PyResult<i32> {
let mut tot = 0_i32;
for d in a.iter() {
tot += d[&key];
}
Ok(tot)
}
如果您想使用 PyList
、PyDict
等,这也适用:
#[pyfunction]
fn sum_list_dicts(a: &PyList,key: String) -> PyResult<i32> {
let mut tot = 0_i32;
for d in a.iter() {
tot += d.downcast::<PyDict>()?.get_item(&key).unwrap().downcast::<PyInt>()?.extract::<i32>()?;
}
Ok(tot)
}
无论使用哪种方式,您都可以简单地从 Python 端调用它:
a = [{"ham":0,"eggs":0},{"eggs": 1},{"eggs": 3,"spam":2}]
b = sum_list_dicts(a,"eggs")
print(b)
>>> 4