Pybind11 默认参数 numpy 数组或 None

Pybind11 default argument numpy array or None

本文关键字：None 数组 numpy 默认 参数 Pybind11      更新时间：2023-10-16

我正在包装一些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?
);

谢谢！

使用 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<py:array>或std::optional<py:array_t<your_custom_type>>修改std::optional<int>即可。

相关文章：