PyO3嵌入Python虚拟环境配置方法

本文档旨在解决在使用 PyO3 将 Python 嵌入 Rust 程序时，遇到的 ModuleNotFoundError 错误，尤其是在使用虚拟环境时。PyO3 默认使用全局 Python 安装，但通过手动初始化 Python 解释器并指定虚拟环境路径，可以确保 Rust 代码正确加载虚拟环境中的 Python 包。本文将提供详细的步骤和代码示例，帮助开发者正确配置 PyO3 和虚拟环境，解决依赖问题。

在使用 PyO3 将 Python 嵌入 Rust 程序时，可能会遇到一个常见问题：即使激活了虚拟环境并在其中安装了所需的 Python 包（例如 pyarrow），Rust 代码仍然无法找到这些包，并抛出 ModuleNotFoundError。这是因为 PyO3 默认情况下会使用全局 Python 安装，而不是虚拟环境。要解决这个问题，需要手动初始化 Python 解释器，并配置正确的虚拟环境路径。

以下是详细的步骤和代码示例：

1. 禁用 auto-initialize 特性

   首先，需要在 Cargo.toml 文件中禁用 pyo3 的 auto-initialize 特性。这将允许我们手动初始化 Python 解释器。

   [dependencies]
   pyo3 = { version = "0.20.0", features = [] } # 移除 "auto-initialize"
   polars = "0.35.4"
   pyo3-polars = "0.9.0"
   libc = "0.2.150"

2. 手动初始化 Python 解释器

   接下来，需要编写 Rust 代码来手动初始化 Python 解释器，并设置虚拟环境的路径。以下是一个示例函数：

   use std::mem::size_of;
   use std::ptr::addr_of_mut;
   
   use libc::wchar_t;
   use pyo3::ffi::*;
   
   fn init_pyo3_with_venv(env_dir: &str) {
       unsafe {
           fn check_exception(status: PyStatus, config: &mut PyConfig) {
               unsafe {
                   if PyStatus_Exception(status) != 0 {
                       PyConfig_Clear(config);
                       if PyStatus_IsExit(status) != 0 {
                           std::process::exit(status.exitcode);
                       }
                       Py_ExitStatusException(status);
                   }
               }
           }
   
           let mut config = std::mem::zeroed::();
           PyConfig_InitPythonConfig(&mut config);
   
           config.install_signal_handlers = 0;
   
           // `wchar_t` is a mess.
           let env_dir_utf16;
           let env_dir_utf32;
           let env_dir_ptr;
           if size_of::() == size_of::() {
               env_dir_utf16 = env_dir
                   .encode_utf16()
                   .chain(std::iter::once(0))
                   .collect::>();
               env_dir_ptr = env_dir_utf16.as_ptr().cast::();
           } else if size_of::() == size_of::() {
               env_dir_utf32 = env_dir
                   .chars()
                   .chain(std::iter::once('\0'))
                   .collect::>();
               env_dir_ptr = env_dir_utf32.as_ptr().cast::();
           } else {
               panic!("unknown encoding for `wchar_t`");
           }
           check_exception(
               PyConfig_SetString(
                   addr_of_mut!(config),
                   addr_of_mut!(config.prefix),
                   env_dir_ptr,
               ),
               &mut config,
           );
   
           check_exception(Py_InitializeFromConfig(&config), &mut config);
   
           PyConfig_Clear(&mut config);
   
           PyEval_SaveThread();
       }
   }

   这个函数使用底层的 C API 来初始化 Python 解释器，并设置 prefix 属性为虚拟环境的路径。

3. 在 main 函数中使用初始化函数

   在 main 函数中，首先获取虚拟环境的路径，然后调用 init_pyo3_with_venv 函数来初始化 Python 解释器。

   use polars::prelude::*;
   use pyo3::{prelude::*, types::PyModule};
   use pyo3_polars::PyDataFrame;
   
   fn main() -> PyResult<()> {
       let env_dir = std::env::current_dir()?.join(".venv");
       if !env_dir.is_dir() {
           panic!("please run from proper directory");
       }
       init_pyo3_with_venv(env_dir.to_str().unwrap());
   
       let code = include_str!("./test.py");
   
       Python::with_gil(|py| {
           let activators = PyModule::from_code(py, code, "activators.py", "activators")?;
   
           let df: DataFrame = df!(
               "integer" => &[1, 2, 3, 4, 5],
               "float" => &[4.0, 5.0, 6.0, 7.0, 8.0],
           )
           .unwrap();
   
           let relu_result: PyDataFrame = activators
               .getattr("test")?
               .call1((PyDataFrame { 0: df },))?
               .extract()?;
   
           Ok(())
       })
   }

   在这个示例中，假设虚拟环境位于项目根目录下的 .venv 目录中。

4. Python 代码示例

   以下是一个简单的 Python 代码示例，用于测试虚拟环境是否配置正确。

   # test.py
   
   def test(x):
       import sys
       print(sys.executable, sys.path, sys.prefix)
   
       import pyarrow
   
       # manipulate dataframe x
       return x

   这个 Python 代码会打印 Python 解释器的路径、模块搜索路径和前缀，并尝试导入 pyarrow 模块。如果一切配置正确，pyarrow 应该能够成功导入。

注意事项：

• 确保在运行 Rust 代码之前，已经激活了虚拟环境，并在其中安装了所需的 Python 包。
• libc 依赖项是必需的，因为它提供了与底层 C API 交互的接口。
• wchar_t 的大小可能因平台而异，因此需要根据实际情况进行处理。
• 此方法绕过了 PyO3 的自动初始化，因此需要手动处理 Python 解释器的初始化和清理。

总结：

通过手动初始化 Python 解释器并配置虚拟环境路径，可以解决在使用 PyO3 嵌入 Python 时遇到的 ModuleNotFoundError 错误。这种方法允许 Rust 代码正确加载虚拟环境中的 Python 包，并确保程序能够正常运行。虽然这种方法需要更多的手动配置，但它提供了更大的灵活性和控制权，特别是在处理复杂的依赖关系时。