使用PyO3向固定签名Rust k近邻函数传递Python嵌套列表的实现

• 原Rust库要求的固定长度切片参数&[[f32; 2]]、&[[f32; 3]]，要求最内层是内存布局连续的定长数组，和PyO3默认推导的可变长Vec<Vec<f32>>类型内存结构、类型定义完全不兼容，无法直接匹配。
• 禁忌操作：不要通过强制转换把&[Vec<f32>]转成目标定长切片，两者内存布局不匹配，会触发未定义行为。

• 调整PyO3函数的入参类型，不直接用Vec<Vec<f32>>做自动类型映射，改为接收通用Python对象&PyAny，逐元素做校验和转换，在Rust端构造出符合原函数要求的定长数组集合。
• 对x参数的处理：逐行遍历Python传入的嵌套列表，校验每个子列表的长度严格等于2，把每个子列表的两个浮点数转成[f32; 2]定长数组，收集成Vec<[f32; 2]>。该Vec的as_slice()方法返回的&[[f32; 2]]类型，可直接匹配原函数的x入参要求。
• 对y（训练集）参数的处理逻辑和x完全一致：逐行校验子列表长度为3，转换为[f32; 3]类型收集成Vec，取切片后即可匹配原函数的y入参要求。
• 转换过程中遇到长度不符合要求的输入，直接抛出明确的Python端ValueError，避免Rust端直接panic。

use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;

// 原库固定签名函数，保持原样不做修改
// pub fn k_nearest_neighbours(k: usize, x: &[[f32; 2]], y: &[[f32; 3]]) -> Vec<Option<f32>>

#[pyfunction]
fn py_k_nearest_neighbours(k: usize, x: &PyAny, y: &PyAny) -> PyResult<Vec<Option<f32>>> {
    // 转换x为长度为2的定长数组集合
    let x_fixed: Vec<[f32; 2]> = x.iter()?
        .map(|row| {
            let row_data: Vec<f32> = row?.extract()?;
            if row_data.len() != 2 {
                return Err(PyValueError::new_err("x参数的每个子列表必须包含恰好2个浮点数"));
            }
            Ok([row_data[0], row_data[1]])
        })
        .collect::<PyResult<_>>()?;

    // 转换y为长度为3的定长数组集合
    let y_fixed: Vec<[f32; 3]> = y.iter()?
        .map(|row| {
            let row_data: Vec<f32> = row?.extract()?;
            if row_data.len() != 3 {
                return Err(PyValueError::new_err("y参数的每个子列表必须包含恰好3个浮点数"));
            }
            Ok([row_data[0], row_data[1], row_data[2]])
        })
        .collect::<PyResult<_>>()?;

    // 直接将切片传入原库函数，类型完全匹配
    Ok(k_nearest_neighbours(k, x_fixed.as_slice(), y_fixed.as_slice()))
}

// 模块注册逻辑保持标准PyO3写法
#[pymodule]
fn kNN_bindings(_py: Python, module: &PyModule) -> PyResult<()> {
    module.add_function(wrap_pyfunction!(py_k_nearest_neighbours, module)?)?;
    Ok(())
}

• 不需要做额外的类型封装，直接传入符合长度要求的嵌套浮点数列表即可，示例：

import kNN_bindings

# 每个子列表含2个浮点数的输入
test_points = [[1.2, 3.4], [5.6, 7.8], [9.0, 1.1]]
# 每个子列表含3个浮点数的训练集
train_points = [[0.1, 0.2, 0.3], [1.1, 2.2, 3.3], [4.4, 5.5, 6.6], [7.7, 8.8, 9.9]]

# 直接调用绑定函数
result = kNN_bindings.py_k_nearest_neighbours(k=2, x=test_points, y=train_points)

• 如果传入的子列表长度不满足要求，会直接抛出清晰的ValueError提示，不会出现晦涩的Rust类型不匹配报错。

内容的提问来源于stack exchange，提问作者Casper S