How do I capture numpy and numpy-like buffer with template arguments using the same function?

[anand-avinash]

I am trying to write functions on C++ side that can accept a numpy array or a custom python buffer that behaves like numpy array. I will pass only 1-D float or double object through these buffer for simplicity. As far as only numpy arrays are concerned, I have the following code to capture them:

auto numpy_lambda = []<typename dfloat>(py::array_t<dfloat> &phi) {
  py::buffer_info phi_info = phi.request();
  dfloat *phi_ptr = reinterpret_cast<dfloat *>(phi_info.ptr);
  my_func(phi_info.size, phi_ptr);
};

PYBIND11_MODULE(my_module, m) {
  m.def("my_func", [](py::array_t<double> &phi) {
    numpy_lambda.template operator()<double>(phi);
  });
  m.def("my_func", [](py::array_t<float> &phi) {
    numpy_lambda.template operator()<float>(phi);
  });
}

This way I can use my_module.my_func(phi) in python where phi has to be a numpy array with dtype either np.float32 or np.float64.

But now, I want to construct my_module.my_func() in a way that can also accept a custom python buffer foo::array_t that has similar buffer.format like numpy arrays. A simplest solution would be to change all py::array_t in C++ side to py::buffer but I would lose template capabilities.

Another solution would be to write a different lambda function for foo::array_t and rely on overloading. While this is suitable for small functions, it would lead a lot of code repetition for large and complex functions.

So I want to understand, what is the appropriate solution in this case?

[anand-avinash]

I found the solution to my query. It was simple actually! The key is to simply change the prototype of the lambda function and rely on the overloading within python module. Finally it would look like this:

auto lambda = []<typename dfloat>(py::buffer &phi) {
  py::buffer_info phi_info = phi.request();
  dfloat *phi_ptr = reinterpret_cast<dfloat *>(phi_info.ptr);
  my_func(phi_info.size, phi_ptr);
};

PYBIND11_MODULE(my_module, m) {
  m.def("my_func", [](py::array_t<double> &phi) {
    lambda.template operator()<double>(phi);
  });
  m.def("my_func", [](py::array_t<float> &phi) {
    lambda.template operator()<float>(phi);
  });

  m.def("my_func", [](foo::array_t<double> &phi) {
    lambda.template operator()<double>(phi);
  });
  m.def("my_func", [](foo::array_t<double> &phi) {
    lambda.template operator()<double>(phi);
  });
}