Got a strange type hint for return value: numpy.ndarray[bool[128, n]]

[Yikai-Liao]

The return type in c++ is define as the Pianoroll below:

#define MIDI_PITCH_NUMS 128
typedef Eigen::Array<bool, MIDI_PITCH_NUMS, Eigen::Dynamic> Pianoroll;

Pianoroll empty_pianoroll(uint16_t quantization = 16) const {
    return Pianoroll::Zero(MIDI_PITCH_NUMS, ceil(this->end_time() * (float) quantization / 4));
};

pybind11-stubgen just generate this strange type hint for my project: numpy.ndarray[bool[128, n]]

And this kind of type is wrong in python.

And here is my project based on pybind11: pyscore

[sizmailov]

The annotation comes from your code, there is no issue on pybind11-stubgen side.

[Yikai-Liao]

The annotation comes from your code, there is no issue on pybind11-stubgen side.

So how can I get a right annotation by changing my code？ I'm just confused about it.

[ringohoffman]

@Yikai-Liao I am seeing this on the project I am working on also. The only dependencies we have in common are pybind11 and eigen.

I see the same type of expression here: libigl-python-bindings, which contains a bunch of eigen code. I'm still not exactly sure how they are getting there though.

[sizmailov]

The annotation is generated by pybind11.

pybind11-stubgen has two built-in options to deal with that, see --help for a bit more details.

pybind11-stubgen [-h]
                 [-o OUTPUT_DIR]
                 [--root-suffix ROOT_SUFFIX]
                 [--ignore-invalid-expressions REGEX]
                 [--ignore-invalid-identifiers REGEX]
                 [--ignore-unresolved-names REGEX]
                 [--ignore-all-errors]
                 [--enum-class-locations REGEX:LOC]
                 [--numpy-array-wrap-with-annotated|   #  <-- first 
                  --numpy-array-remove-parameters]      #  <-- second
                 [--print-invalid-expressions-as-is]
                 [--print-safe-value-reprs REGEX]
                 [--exit-code]
                 [--stub-extension EXT]
                 MODULE_NAME

[ringohoffman]

In its type stubs, numpy.ndarray is generic w.r.t. 2 type variables:

from __future__ import annotations

from typing import Any, Generic, TypeVar

import numpy as np

_DType_co = TypeVar("_DType_co", covariant=True, bound=np.dtype[Any])
_ShapeType = TypeVar("_ShapeType", bound=Any)

class ndarray(Generic[_ShapeType, _DType_co]):
    ...

And you can utilize them like this:

from __future__ import annotations

from typing import TypeVar

import numpy as np
from typing_extensions import Literal

M = TypeVar("M")
N = TypeVar("N")
P = TypeVar("P")

def matmul(
    l: np.ndarray[tuple[M, N], np.dtype[np.float64]],
    r: np.ndarray[tuple[N, P], np.dtype[np.float64]],
) -> np.ndarray[tuple[M, P], np.dtype[np.float64]]:
    return l @ r

arr1: np.ndarray[tuple[Literal[2], Literal[3]], np.dtype[np.float64]] = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float64)
arr2: np.ndarray[tuple[Literal[3], Literal[4]], np.dtype[np.float64]] = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=np.float64)

arr3 = matmul(arr1, arr2)  # pyright reveals the type is: ndarray[tuple[Literal[2], Literal[4]], dtype[float64]]
arr3.shape  # (2, 4)

What would you think about generating this style of annotations? You can see that it works, and it seems like this is what people are currently doing in numpy:

numpy/numpy#16544 (comment)
numpy/numpy#16544 (comment)

The pybind11 style:

numpy.ndarray[numpy.float32[m, 1]]

would become:

M = typing.TypeVar("M")
...  # any others

numpy.ndarray[tuple[M, typing_extensions.Literal[1]], numpy.dtype[numpy.float32]]

[sizmailov]

The --numpy-array-wrap-with-annotated currently produces the following stub

pybind11-stubgen/tests/stubs/python-3.12/pybind11-master/demo/_bindings/eigen.pyi

Lines 35 to 38 in afc508f

	def dense_matrix_c(
	arg0: typing.Annotated[
	numpy.ndarray, numpy.float32, pybind11_stubgen.typing_ext.DynamicSize("m", "n")
	]

I think it's good enough.

If you are interested in adding different flavors of annotating numpy arrays, take a look at #113 and #115.