Javier
Javier

Reputation: 341

Pybind11 default argument numpy array or None

I am wrapping some C++ code to use it from Python. I want to call a C++ function with an argument that can take a None value or a numpy.array of the same size of another input variable. This is the example:

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)

How can I code it in C++ with pybind11? I have the following function signatures and code:

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?
    );

Thank you very much!

Upvotes: 3

Views: 3836

Answers (1)

Yixing Lao
Yixing Lao

Reputation: 1428

With C++17's std::optional, here's an example that should work. For earlier versions of C++, you might want a backported optional.h and implement your own optional_caster similar to the one in pybind11/stl.h.

Say you want this function:

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

Here's the equivalent C++ pybind implementation:

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()
);

In python, this function can be called with:

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

For numpy arrays, simply modify the std::optional<int> with std::optional<py:array> or std::optional<py:array_t<your_custom_type>> in the example.

Upvotes: 6

Related Questions