Append numpy prefix to dtype (#179)

Author: sizmailov

pybind11_stubgen/__init__.py

@@ -35,6 +35,7 @@
     FixNumpyArrayDimAnnotation,
     FixNumpyArrayFlags,
     FixNumpyArrayRemoveParameters,
+    FixNumpyDtype,
     FixPEP585CollectionNames,
     FixPybind11EnumStrDoc,
     FixRedundantBuiltinsAnnotation,
@@ -231,6 +232,7 @@ class Parser(
         FixMissingEnumMembersAnnotation,
         OverridePrintSafeValues,
         *numpy_fixes,  # type: ignore[misc]
+        FixNumpyDtype,
         FixNumpyArrayFlags,
         FixCurrentModulePrefixInTypeNames,
         FixBuiltinTypes,

pybind11_stubgen/parser/mixins/fix.py

@@ -586,6 +586,26 @@ def parse_annotation_str(
         return result
 
 
+class FixNumpyDtype(IParser):
+    __numpy_dtype = QualifiedName.from_str("numpy.dtype")
+
+    def parse_annotation_str(
+        self, annotation_str: str
+    ) -> ResolvedType | InvalidExpression | Value:
+        result = super().parse_annotation_str(annotation_str)
+        if (
+            not isinstance(result, ResolvedType)
+            or len(result.name) != 1
+            or result.parameters is not None
+        ):
+            return result
+
+        word = result.name[0]
+        if word != Identifier("dtype"):
+            return result
+        return ResolvedType(name=self.__numpy_dtype)
+
+
 class FixNumpyArrayFlags(IParser):
     __ndarray_name = QualifiedName.from_str("numpy.ndarray")
     __flags: set[QualifiedName] = {

tests/py-demo/bindings/src/modules/numpy.cpp

@@ -9,4 +9,5 @@ void bind_numpy_module(py::module&&m) {
         m.def("accept_ndarray_int", [](py::array_t<int> &) {});
         m.def("accept_ndarray_float64", [](py::array_t<double> &) {});
     }
+    m.def("return_dtype", []() { return py::dtype("<i4"); });
 }

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

@@ -9,6 +9,7 @@ __all__ = [
     "accept_ndarray_int",
     "get_ndarray_float64",
     "get_ndarray_int",
+    "return_dtype",
 ]
 
 def accept_ndarray_float64(
@@ -17,3 +18,4 @@ def accept_ndarray_float64(
 def accept_ndarray_int(arg0: typing.Annotated[numpy.ndarray, numpy.int32]) -> None: ...
 def get_ndarray_float64() -> typing.Annotated[numpy.ndarray, numpy.float64]: ...
 def get_ndarray_int() -> typing.Annotated[numpy.ndarray, numpy.int32]: ...
+def return_dtype() -> numpy.dtype: ...

tests/stubs/python-3.12/pybind11-v2.11/demo/_bindings/numpy.pyi

@@ -9,6 +9,7 @@ __all__ = [
     "accept_ndarray_int",
     "get_ndarray_float64",
     "get_ndarray_int",
+    "return_dtype",
 ]
 
 def accept_ndarray_float64(
@@ -17,3 +18,4 @@ def accept_ndarray_float64(
 def accept_ndarray_int(arg0: typing.Annotated[numpy.ndarray, numpy.int32]) -> None: ...
 def get_ndarray_float64() -> typing.Annotated[numpy.ndarray, numpy.float64]: ...
 def get_ndarray_int() -> typing.Annotated[numpy.ndarray, numpy.int32]: ...
+def return_dtype() -> numpy.dtype: ...

tests/stubs/python-3.12/pybind11-v2.9/demo/_bindings/numpy.pyi

@@ -9,6 +9,7 @@ __all__ = [
     "accept_ndarray_int",
     "get_ndarray_float64",
     "get_ndarray_int",
+    "return_dtype",
 ]
 
 def accept_ndarray_float64(
@@ -17,3 +18,4 @@ def accept_ndarray_float64(
 def accept_ndarray_int(arg0: typing.Annotated[numpy.ndarray, numpy.int32]) -> None: ...
 def get_ndarray_float64() -> typing.Annotated[numpy.ndarray, numpy.float64]: ...
 def get_ndarray_int() -> typing.Annotated[numpy.ndarray, numpy.int32]: ...
+def return_dtype() -> numpy.dtype: ...

tests/stubs/python-3.7/pybind11-master/demo/_bindings/numpy.pyi

@@ -8,6 +8,7 @@ __all__ = [
     "accept_ndarray_int",
     "get_ndarray_float64",
     "get_ndarray_int",
+    "return_dtype",
 ]
 
 def accept_ndarray_float64(
@@ -20,3 +21,4 @@ def get_ndarray_float64() -> typing_extensions.Annotated[
     numpy.ndarray, numpy.float64
 ]: ...
 def get_ndarray_int() -> typing_extensions.Annotated[numpy.ndarray, numpy.int32]: ...
+def return_dtype() -> numpy.dtype: ...

tests/stubs/python-3.8/pybind11-master/demo/_bindings/numpy.pyi

@@ -8,6 +8,7 @@ __all__ = [
     "accept_ndarray_int",
     "get_ndarray_float64",
     "get_ndarray_int",
+    "return_dtype",
 ]
 
 def accept_ndarray_float64(
@@ -20,3 +21,4 @@ def get_ndarray_float64() -> typing_extensions.Annotated[
     numpy.ndarray, numpy.float64
 ]: ...
 def get_ndarray_int() -> typing_extensions.Annotated[numpy.ndarray, numpy.int32]: ...
+def return_dtype() -> numpy.dtype: ...