使用pyo3时如何在Rust程序中捕获Python解释器的stdout输出

实现原理 #

我们通过临时替换Python的sys.stdout对象为自定义实现，将所有写入stdout的调用转发到Rust侧处理，执行脚本完成后再恢复原始stdout，即可实现输出拦截，不会写入真实终端。

前置依赖配置 #

首先在Cargo.toml中添加依赖：

[package]
name = "pyo3-capture-demo"
version = "0.1.0"
edition = "2021"

[dependencies]
pyo3 = { version = "0.20", features = ["auto-initialize"] }

方案1：全量捕获输出到Rust字符串 #

适合需要等脚本执行完一次性获取所有输出的场景：

use pyo3::prelude::*;

fn run_capture_output(script: &str) -> PyResult<String> {
    Python::with_gil(|py| {
        // 保存原始stdout用于后续恢复
        let sys = py.import("sys")?;
        let original_stdout = sys.getattr("stdout")?;

        // 定义Python侧的自定义stdout类，write方法将内容写入缓冲区
        let stdout_class = py.eval(r#"
class CapturedStdout:
    def __init__(self):
        self.buffer = ""
    def write(self, s):
        self.buffer += s
    def flush(self):
        pass
"#, None, None)?.cast_as::<PyAny>()?;

        // 实例化自定义stdout并替换系统stdout
        let captured_stdout = stdout_class.call0()?;
        sys.setattr("stdout", captured_stdout)?;

        // 执行目标脚本，捕获可能的异常保证后续能恢复stdout
        let run_result = py.run(script, None, None);

        // 恢复原始stdout
        sys.setattr("stdout", original_stdout)?;

        // 抛出执行异常（如果有）
        run_result?;

        // 读取捕获的输出返回
        let captured = captured_stdout.getattr("buffer")?.extract::<String>()?;
        Ok(captured)
    })
}

fn main() -> PyResult<()> {
    let output = run_capture_output("print('abc')")?;
    // 这里的output就是捕获到的"abc\n"，不会打印到终端
    assert_eq!(output, "abc\n");
    println!("Rust侧拿到的输出：{}", output.trim());
    Ok(())
}

方案2：逐行回调实时处理输出 #

适合需要边执行边处理输出的场景，输出产生时会逐行传递给Rust回调函数：

use pyo3::prelude::*;
use std::cell::RefCell;

fn run_with_stdout_callback<F>(script: &str, mut callback: F) -> PyResult<()>
where
    F: FnMut(&str) + 'static,
{
    // 用RefCell包装回调，方便在Python方法中调用
    let callback_cell = RefCell::new(callback);

    Python::with_gil(|py| {
        let sys = py.import("sys")?;
        let original_stdout = sys.getattr("stdout")?;

        // 定义自定义stdout类，将write收到的内容按换行分割后调用Rust回调
        #[pyclass]
        struct CallbackStdout {
            buffer: String,
            callback: Box<dyn FnMut(&str)>,
        }

        #[pymethods]
        impl CallbackStdout {
            fn write(&mut self, s: &str) {
                self.buffer.push_str(s);
                // 按换行分割，逐行调用回调
                while let Some(newline_pos) = self.buffer.find('\n') {
                    let line = self.buffer.drain(..newline_pos + 1).collect::<String>();
                    (self.callback)(&line);
                }
            }

            fn flush(&self) {}
        }

        // 实例化自定义stdout
        let captured_stdout = Py::new(py, CallbackStdout {
            buffer: String::new(),
            callback: Box::new(move |line| callback_cell.borrow_mut()(line)),
        })?;

        sys.setattr("stdout", captured_stdout.as_ref(py))?;

        // 执行脚本+恢复stdout逻辑同方案1
        let run_result = py.run(script, None, None);
        sys.setattr("stdout", original_stdout)?;
        run_result?;

        Ok(())
    })
}

fn main() -> PyResult<()> {
    run_with_stdout_callback("print('line1'); print('line2'); import time; time.sleep(1); print('line3')", |line| {
        // 每行输出产生时会立刻调用这个闭包
        println!("Rust收到新行：{}", line.trim());
    })?;
    Ok(())
}

补充说明 #

• 如果需要同时捕获stderr，按照相同逻辑替换sys.stderr即可
• 如果执行的Python脚本中存在手动写入sys.__stdout__的逻辑，这种方案无法拦截，属于Python标准行为

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