如何通过PyBind11直接返回Numpy数组，无需Python类型转换？

完全可以通过PyBind11实现需求，根据是否需要共享内存，有两种主流实现方式：

方式一：零拷贝直接包装缓冲区（推荐） #

利用PyBind11的py::array类，直接基于你的C++缓冲区类构造numpy数组，无需拷贝数据，同时确保内存安全。

实现代码示例 #

假设你的C++缓冲区类MyBufferClass已实现data()（返回数据指针）、shape()（返回维度数组）、strides()（返回字节步长数组）方法：

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

namespace py = pybind11;

class MyBufferClass {
public:
    MyBufferClass(size_t rows, size_t cols) : m_rows(rows), m_cols(cols) {
        m_data.resize(rows * cols);
        // 填充示例数据
        for (size_t i = 0; i < rows * cols; ++i) {
            m_data[i] = static_cast<float>(i);
        }
    }

    float* data() { return m_data.data(); }
    const float* data() const { return m_data.data(); }

    std::vector<ssize_t> shape() const {
        return {static_cast<ssize_t>(m_rows), static_cast<ssize_t>(m_cols)};
    }

    std::vector<ssize_t> strides() const {
        // 按numpy行优先布局计算字节步长
        return {
            static_cast<ssize_t>(sizeof(float) * m_cols),
            static_cast<ssize_t>(sizeof(float))
        };
    }

private:
    std::vector<float> m_data;
    size_t m_rows, m_cols;
};

// 直接返回numpy数组的函数
py::array get_numpy_array() {
    MyBufferClass buf(3, 4); // 创建缓冲区实例
    // 构造numpy数组，绑定缓冲区的生命周期
    return py::array(
        py::dtype::of<float>(),
        buf.shape(),
        buf.strides(),
        buf.data(),
        py::cast(buf) // 将C++实例作为numpy数组的所有者，避免悬空指针
    );
}

PYBIND11_MODULE(my_buffer_module, m) {
    m.def("get_numpy_array", &get_numpy_array, "直接返回基于C++缓冲区的Numpy数组");
}

关键说明 #

• 最后一个参数py::cast(buf)是核心：它将C缓冲区实例转为Python对象，作为numpy数组的"所有者"。当Python端的numpy数组被垃圾回收时，C实例才会被销毁，确保数据指针始终有效。
• 步长必须以字节为单位，而非元素个数，要和numpy的内存布局匹配。

方式二：拷贝数据到新的Numpy数组 #

如果不需要共享内存（比如C++缓冲区会被立即销毁），可以先创建空的numpy数组，再拷贝数据：

py::array copy_to_numpy() {
    MyBufferClass buf(3, 4);
    auto shape = buf.shape();
    // 创建空的numpy数组
    py::array arr(py::dtype::of<float>(), shape);
    // 拷贝数据到numpy数组
    std::memcpy(
        arr.mutable_data(),
        buf.data(),
        shape[0] * shape[1] * sizeof(float)
    );
    return arr;
}

注意事项 #

• 零拷贝方式下，确保C++缓冲区的内存不会在numpy数组存活期间被释放。
• 如果你的C++类使用自定义内存分配（而非std::vector），需要确保内存释放逻辑正确绑定到Python的垃圾回收机制。

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