[OTHER]: Generated stubs for nb::ndarray are too strict about dtype (e.g., NDArray[float64]) when nanobind converts inputs implicitly

[nurpax]

Summary

Recent nanobind generates .pyi where parameters typed as nb::ndarray<const double, ...> become Annotated[NDArray[float64], ...]. Static checkers (e.g., VS Code Pylance) then reject valid calls that pass non-float64 arrays -- even though nanobind converts these to double at runtime. Older nanobind emitted ArrayLike, which didn’t trigger false positives.

Request: a way to generate looser stubs to reflect accepted inputs when dtype conversion is allowed.

---

Environment

• nanobind @ 10cc510
• Affected checker: VS Code Pylance (strict about NDArray[...] vs ArrayLike)

Background / behavior change

Example C++ binding:

using ndarray_1d = nb::ndarray<const double, nb::ndim<1>, nb::device::cpu, nb::c_contig>;

m.def("plot_line", [](const char* label_id, ndarray_1d& xs, ndarray_1d& ys, ImPlotLineFlags_ flags) {
    ImPlot::PlotLine(label_id, (const double*)xs.data(), (const double*)ys.data(), xs.shape(0), flags);
}, "label_id"_a, "xs"_a, "ys"_a, "flags"_a.sig("LineFlags.NONE") = ImPlotLineFlags_None);

Old stubs:

xs: Annotated[ArrayLike, dict(dtype='float64', shape=(None), order='C', device='cpu', writable=False)]

New stubs that issue a type error on valid inputs:

xs: Annotated[NDArray[float64], dict(shape=(None,), order='C', device='cpu', writable=False)]

Problem (false positive type errors)

A call like this is valid at runtime (nanobind converts to double) but flagged by Pylance:

import numpy as np

data_x = np.arange(0, 10, 0.25)      # NDArray[floating[Any]]
data_y = np.sin(data_x)
implot.plot_line("Sine Wave", data_x, data_y)

Error:

Argument of type "_Array1D[floating[Any]]" cannot be assigned to parameter "xs" of type "NDArray[float64]" ...

Why this seems incorrect

The C++ signature allows implicit conversion to double (const double element type). nanobind does the conversion automatically.  The generated stub, however, requires NDArray[np.float64], which is stricter than the actual accepted inputs.

Attempted workaround (dtype-agnostic typedef)

Changing the binding to a dtype-agnostic ndarray and converting inside:

using ndarray_1d_any = nb::ndarray<nb::ndim<1>, nb::device::cpu, nb::c_contig>; // no double constraint

m.def("plot_line", [](const char* label_id, ndarray_1d_any& xs, ndarray_1d_any& ys, ImPlotLineFlags_ flags) {
    auto xv = xs.view<const double, nb::ndim<1>, nb::device::cpu, nb::c_contig>();
    auto yv = ys.view<const double, nb::ndim<1>, nb::device::cpu, nb::c_contig>();
    ImPlot::PlotLine(label_id, (const double*)xv.data(), (const double*)yv.data(), xv.shape(0), flags);
}, "label_id"_a, "xs"_a, "ys"_a, "flags"_a.sig("LineFlags.NONE") = ImPlotLineFlags_None);

This produces the desired looser stub:

xs: Annotated[NDArray, dict(shape=(None,), order='C', device='cpu')]

…but the .view<...>() does not perform dtype conversion, so the implementation isn’t correct for non-float64 inputs.

Proposal directions

1. Stub generation policy knob: emit NDArray[Any] (or ArrayLike) for convertible input arguments unless .noconvert() is present.

2. Fine-grained overrides: allow "arg"_a.sig("NDArray[Any]") to override type of a single argument without re-specifying the whole signature.

3. Declare bindings without nb::ndarray<> constraints, allow binding code to perform explicit array conversion inside the binding code.  How to do  this currently -- this would be a decent workaround?

Workarounds today

• Explicit Python-side conversion (np.asarray(x, dtype=np.float64, order='C')).
• Full manual stub override with nb::sig(...) (tedious for multi-arg functions).
• Stick with older nanobind (where ArrayLike was used).
• Post-process the generated .pyi files to sed NDArray[np.float64] -> NDArray[Any].

[nurpax]

Previously discussed here: #1152

The new style stub gen was requested in a number of issues: #442, #1094, #983, and #984.

Paging in authors/commenters from the above issues, in case someone more knowledgeable about Python stubs for NDArrays would have good ideas for refining stubgen: @hpkfft, @yosh-matsuda, @breathe, @WKarel, @Trzeth, @pwuertz, @keltonhalbert, @AaronSartin.

[keltonhalbert]

My use case with the new stub style was primarily focused on getting the accurate numpy.ndarray return type in order to get quality autocomplete and type info from the LSP knowing it is, in fact, a numpy.ndarray. As far as the input arguments, I was surprised to see those changed from ArrayLike since I thought I read in another PR/issue about ArrayLike being used for this exact reason.

At the very least, some sort of flag seems like a reasonable solution. I also would not be opposed to returning to ArrayLike for input stubs outright. I'm not sure if there are qualms about differing signatures for input and output, but in my case, I am being as strongly typed as possible for my output:

namespace nb = nanobind;

// read only input arrays
using const_prof_arr_t =
    nb::ndarray<float, nb::ndim<1>, nb::device::cpu, nb::c_contig, nb::ro>;

// read/write arrays
using prof_arr_t =
    nb::ndarray<float, nb::ndim<1>, nb::device::cpu, nb::c_contig>;

using out_arr_t = nb::ndarray<nb::numpy, float, nb::ndim<1>, nb::c_contig>;

Edit: I think the reason I didn't notice this either, is I strongly type my input arrays python-side with astype='float32'

[breathe]

Previously discussed here: #1152

The new style stub gen was requested in a number of issues: #442, #1094, #983, and #984.

Paging in authors/commenters from the above issues, in case someone more knowledgeable about Python stubs for NDArrays would have good ideas for refining stubgen: @hpkfft, @yosh-matsuda, @breathe, @WKarel, @Trzeth, @pwuertz, @keltonhalbert, @AaronSartin.

I just updated our project to the latest release to test the new stupgen logic -- and the new setup is an improvement in my codebase ... -- however I am using Eigen:: types on the cpp side rather than ndarray:: directly

Our codebase previously needed to use:

# HACK HACK HACK: hopefully temporary -- see this github discussion -- we patch the generated .pyi file: https://github.com/wjakob/nanobind/discussions/630
install(CODE "execute_process(COMMAND /bin/sh -c \"sed -i.bak 's/ArrayLike/NDArray/g' '${PYI_FILE}' && rm '${PYI_FILE}.bak' \")")

Because ArrayLike is far to accepting a type. You can pass pretty much anything to it -- including many more varieties of python values than the underlying cpp functions will accept. With ArrayLike typing -- many python values passed to the cpp function will fail to produce type errors and running the code will lead to nanobind being forced to issue a runtime type conversion error...

With mypy 1.9.0 at least -- I'm not actually repro'ing your failure to typecheck passing values of type ndarray[floating[any]] to functions expecting NDArray[np.float64]

I'm using Eigen:: instead of ndarray on the cpp side -- but the generated pyi types seem to match ... and your error should depend on the .pyi file only afaict

auto test_function(const Eigen::VectorXd &x) -> Eigen::VectorXd
{
    return x;
}

NB_MODULE(zephyrus_native_ext, m)
{
    m.def("test_function", &test_function, nb::arg("x"));
...

.pyi

def test_function(x: Annotated[NDArray[numpy.float64], dict(shape=(None,), order='C')]) -> Annotated[NDArray[numpy.float64], dict(shape=(None,), order='C')]: ...

from zephyrus.lib.native import test_function

data_x = np.arange(0, 10, 0.25)
test_function(data_x)

mypy thinks this is fine and doesn't issue the type error you see about being unable to convert ndarray[floating[any]] to NDArray[float64]

Question: Is the type checking error actually resulting from the ndarray floating portion of this or does type checker behavior change vs what I observe here -- as a result of the extra annotated properties possibly ?
Answer: no I don't think so -- I manually added device='cpu', writable=False to my .pyi file and it still typechecks for me ...

Is it possible we are observing a different type-checking semantic in mypy vs pylance?

Seems like there could also be numpy version differences involved ... It appears that numpy may have changed behavior here in a way that would affect this example ... numpy/numpy#27957 (comment) ... I'm on an old numpy numpy = "1.26.4"

[hpkfft]

For reference: JAX Typing Best Practices