如何在PyO3 Rust Python扩展中保存numpy.ndarray至静态/线程局部变量

问题背景 #

你尝试用rust-numpy crate构建Python扩展，希望将numpy.ndarray对应的PyReadwriteArray1<i32>保存到线程局部变量中供后续处理，但遇到了生命周期编译错误：

static ARRAYS: RefCell<Vec::<PyReadwriteArray1<i32>>> = RefCell::new(Vec::<PyReadwriteArray1<i32>>::new());
  |                                                   ^ expected named lifetime parameter

核心原因 #

PyReadwriteArray1<'py, T>是带有生命周期的借用类型，它绑定到Python GIL的生命周期'py，只能在当前GIL作用域内有效。而静态/线程局部变量的生命周期是'static，无法存储带有临时生命周期的借用对象——这违反了Rust的生命周期安全规则。

但你完全可以在Rust扩展中保存Python对象，只是需要用PyO3提供的全局引用类型来存储，而不是借用类型。

解决方案 #

使用PyO3的Py<T>类型存储Python对象的全局引用，它的生命周期是'static，可以安全地存到线程局部变量中。具体步骤如下：

1. 将PyReadwriteArray1转换为Py<PyArray1<i32>>（全局引用），需要借助Python GIL环境
2. 线程局部变量存储Vec<Py<PyArray1<i32>>>
3. 后续需要操作数组时，再将全局引用转换回可读写的数组视图（需在GIL作用域内）

修正后的代码：

use std::cell::RefCell;
use numpy::{PyArray1, PyReadwriteArray1};
use pyo3::{Bound, Py, PyResult, Python, pymodule, types::PyModule};

// 线程局部变量存储numpy数组的全局引用
thread_local! {
    static ARRAYS: RefCell<Vec<Py<PyArray1<i32>>>> = RefCell::new(Vec::new());
}

#[pymodule]
fn rust_ext<'py>(m: &Bound<'py, PyModule>) -> PyResult<()> {
    #[pyfn(m)]
    fn register_array(x: PyReadwriteArray1<i32>, py: Python<'_>) {
        // 修改数组内容
        x.as_array_mut()[0] = 100;
        
        // 将PyReadwriteArray1转换为全局引用Py<PyArray1<i32>>
        let py_array: Py<PyArray1<i32>> = x.into_py(py);
        
        // 存入线程局部变量
        ARRAYS.with_borrow_mut(|v| v.push(py_array));
    }

    // 示例：从线程局部变量取出并操作数组
    #[pyfn(m)]
    fn process_stored_arrays(py: Python<'_>) -> PyResult<()> {
        ARRAYS.with_borrow_mut(|v| {
            for py_array in v.iter() {
                // 将全局引用转换为可读写视图
                let mut array = PyReadwriteArray1::from_bound(py_array.bind(py));
                array.as_array_mut()[0] += 10;
            }
        });
        Ok(())
    }

    Ok(())
}

关键细节说明 #

• Py<T>：PyO3提供的全局引用类型，持有Python对象的所有权，会自动参与Python的垃圾回收，生命周期为'static，可以安全存储到静态/线程局部变量中
• into_py(py)：将PyReadwriteArray1转换为Py<PyArray1>的方法，需要传入当前GIL环境Python<'_>
• from_bound：后续操作时，通过py_array.bind(py)将全局引用绑定到当前GIL，再转换为PyReadwriteArray1进行读写操作

注意事项 #

• 线程安全：线程局部变量确保每个线程的数组存储相互独立，符合Python GIL的线程模型
• 内存管理：避免循环引用，若你的扩展对象和存储的数组形成循环引用，需要手动打破（比如使用弱引用）
• GIL持有：所有操作Python对象的代码必须在GIL作用域内，pyfn会自动获取GIL，手动操作时需要用Python::with_gil包裹

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