Pybind11 默认参数 numpy 数组或无

Posted 2023-02-23

【中文标题】Pybind11 默认参数 numpy 数组或无

【英文标题】：Pybind11 default argument numpy array or None

【发布时间】：2019-12-10 15:29:13

【问题描述】：

我正在包装一些 C++ 代码以在 Python 中使用它。我想调用一个带有参数的 C++ 函数，该参数可以采用与另一个输入变量相同大小的 None 值或 numpy.array。这是一个例子：

import example

# Let I a numpy array containing a 2D or 3D image
M = I > 0
# Calling the C++ function with no mask
example.fit(I, k, mask=None)
# Calling the C++ function with mask as a Numpy array of the same size of I
example.fit(I, k, mask=M)

如何使用 pybind11 在 C++ 中对其进行编码？我有以下函数签名和代码：

void fit(const py::array_t<float, py::array::c_style | py::array::forcecast> &input, 
         int k,
         const py::array_t<bool, py::array::c_style | py::array::forcecast> &mask)

    ...


PYBIND11_MODULE(example, m)

    m.def("fit", &fit,
        py::arg("input"),
        py::arg("k"),
        py::arg("mask") = nullptr // Don't know what to put here?
    );

非常感谢！

【问题讨论】：

我认为您可以使用与this question相同的解决方案

【参考方案1】：

使用 C++17 的 std::optional，这是一个应该可以工作的示例。对于早期版本的 C++，您可能需要一个向后移植的 optional.h 并实现您自己的 optional_caster，类似于 pybind11/stl.h 中的那个。

说你想要这个功能：

def add(a, b=None):
    # Assuming a, b are int.
    if b is None:
        return a
    else:
        return a + b

这是等效的 C++ pybind 实现：

m.def("add",
    [](int a, std::optional<int> b) 
        if (!b.has_value()) 
            return a;
         else 
            return a + b.value();
        
    ,
    py::arg("a"), py::arg("b") = py::none()
);

在python中，这个函数可以调用：

add(1)
add(1, 2)
add(1, b=2)
add(1, b=None)

对于 numpy 数组，只需在示例中将 std::optional<int> 修改为 std::optional<py:array> 或 std::optional<py:array_t<your_custom_type>>。