ENH: make dtype generic over scalar type

This allows representing dtype subclasses via constructs like
`np.dtype[np.float64]`.

5 files changed, 151 insertions, 22 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 3d40682e7..139f2a1bc 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -15,6 +15,8 @@ from numpy.typing import (
     _FloatLike,
     _ComplexLike,
     _NumberLike,
+    _SupportsDtype,
+    _VoidDtypeLike,
 )
 from numpy.typing._callable import (
     _BoolOp,
@@ -527,16 +529,112 @@ where: Any
 who: Any
 
 _NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray)
+_DTypeScalar = TypeVar("_DTypeScalar", bound=generic)
 _ByteOrder = Literal["S", "<", ">", "=", "|", "L", "B", "N", "I"]
 
-class dtype:
+class dtype(Generic[_DTypeScalar]):
     names: Optional[Tuple[str, ...]]
-    def __init__(
-        self,
-        dtype: DtypeLike,
+    # Overload for subclass of generic
+    @overload
+    def __new__(
+        cls,
+        dtype: Type[_DTypeScalar],
         align: bool = ...,
         copy: bool = ...,
-    ) -> None: ...
+    ) -> dtype[_DTypeScalar]: ...
+    # Overloads for string aliases
+    @overload
+    def __new__(
+        cls,
+        dtype: Literal["float64", "f8", "<f8", ">f8", "float", "double", "float_", "d"],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[float64]: ...
+    @overload
+    def __new__(
+        cls,
+        dtype: Literal["float32", "f4", "<f4", ">f4", "single"],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[float32]: ...
+    @overload
+    def __new__(
+        cls,
+        dtype: Literal["int64", "i8", "<i8", ">i8"],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[int64]: ...
+    @overload
+    def __new__(
+        cls,
+        dtype: Literal["int32", "i4", "<i4", ">i4"],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[int32]: ...
+    # "int" resolves to int_, which is system dependent, and as of now
+    # untyped. Long-term we'll do something fancier here.
+    @overload
+    def __new__(
+        cls,
+        dtype: Literal["int"],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype: ...
+    # Overloads for Python types. Order is important here.
+    @overload
+    def __new__(
+        cls,
+        dtype: Type[bool],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[bool_]: ...
+    # See the notes for "int"
+    @overload
+    def __new__(
+        cls,
+        dtype: Type[int],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[Any]: ...
+    @overload
+    def __new__(
+        cls,
+        dtype: Type[float],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[float64]: ...
+    # None is a special case
+    @overload
+    def __new__(
+        cls,
+        dtype: None,
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[float64]: ...
+    # dtype of a dtype is the same dtype
+    @overload
+    def __new__(
+        cls,
+        dtype: dtype[_DTypeScalar],
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[_DTypeScalar]: ...
+    # TODO: handle _SupportsDtype better
+    @overload
+    def __new__(
+        cls,
+        dtype: _SupportsDtype,
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[Any]: ...
+    # Catchall overload
+    @overload
+    def __new__(
+        cls,
+        dtype: _VoidDtypeLike,
+        align: bool = ...,
+        copy: bool = ...,
+    ) -> dtype[void]: ...
     def __eq__(self, other: DtypeLike) -> bool: ...
     def __ne__(self, other: DtypeLike) -> bool: ...
     def __gt__(self, other: DtypeLike) -> bool: ...
diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
index 987aa39aa..dafabd95a 100644
--- a/numpy/typing/__init__.py
+++ b/numpy/typing/__init__.py
@@ -102,7 +102,7 @@ from ._scalars import (
 )
 from ._array_like import _SupportsArray, ArrayLike
 from ._shape import _Shape, _ShapeLike
-from ._dtype_like import DtypeLike
+from ._dtype_like import _SupportsDtype, _VoidDtypeLike, DtypeLike
 
 from numpy._pytesttester import PytestTester
 test = PytestTester(__name__)
diff --git a/numpy/typing/_dtype_like.py b/numpy/typing/_dtype_like.py
index 7c1946a3e..5bfd8ffdc 100644
--- a/numpy/typing/_dtype_like.py
+++ b/numpy/typing/_dtype_like.py
@@ -38,18 +38,9 @@ else:
     _DtypeDict = Any
     _SupportsDtype = Any
 
-# Anything that can be coerced into numpy.dtype.
-# Reference: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
-DtypeLike = Union[
-    dtype,
-    # default data type (float64)
-    None,
-    # array-scalar types and generic types
-    type,  # TODO: enumerate these when we add type hints for numpy scalars
-    # anything with a dtype attribute
-    _SupportsDtype,
-    # character codes, type strings or comma-separated fields, e.g., 'float64'
-    str,
+
+# Would create a dtype[np.void]
+_VoidDtypeLike = Union[
     # (flexible_dtype, itemsize)
     Tuple[_DtypeLikeNested, int],
     # (fixed_dtype, shape)
@@ -67,6 +58,21 @@ DtypeLike = Union[
     Tuple[_DtypeLikeNested, _DtypeLikeNested],
 ]
 
+# Anything that can be coerced into numpy.dtype.
+# Reference: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
+DtypeLike = Union[
+    dtype,
+    # default data type (float64)
+    None,
+    # array-scalar types and generic types
+    type,  # TODO: enumerate these when we add type hints for numpy scalars
+    # anything with a dtype attribute
+    _SupportsDtype,
+    # character codes, type strings or comma-separated fields, e.g., 'float64'
+    str,
+    _VoidDtypeLike,
+]
+
 # NOTE: while it is possible to provide the dtype as a dict of
 # dtype-like objects (e.g. `{'field1': ..., 'field2': ..., ...}`),
 # this syntax is officially discourged and
diff --git a/numpy/typing/tests/data/fail/dtype.py b/numpy/typing/tests/data/fail/dtype.py
index 3dc027daf..7d4783d8f 100644
--- a/numpy/typing/tests/data/fail/dtype.py
+++ b/numpy/typing/tests/data/fail/dtype.py
@@ -1,15 +1,16 @@
 import numpy as np
 
-
 class Test:
     not_dtype = float
 
 
-np.dtype(Test())  # E: Argument 1 to "dtype" has incompatible type
+np.dtype(Test())  # E: No overload variant of "dtype" matches
 
-np.dtype(
-    {  # E: Argument 1 to "dtype" has incompatible type
+np.dtype(  # E: No overload variant of "dtype" matches
+    {
         "field1": (float, 1),
         "field2": (int, 3),
     }
 )
+
+np.dtype[np.float64](np.int64)  # E: Argument 1 to "dtype" has incompatible type
diff --git a/numpy/typing/tests/data/reveal/dtype.py b/numpy/typing/tests/data/reveal/dtype.py
new file mode 100644
index 000000000..aca7e8a5e
--- /dev/null
+++ b/numpy/typing/tests/data/reveal/dtype.py
@@ -0,0 +1,24 @@
+import numpy as np
+
+reveal_type(np.dtype(np.float64))  # E: numpy.dtype[numpy.float64*]
+reveal_type(np.dtype(np.int64))  # E: numpy.dtype[numpy.int64*]
+
+# String aliases
+reveal_type(np.dtype("float64"))  # E: numpy.dtype[numpy.float64]
+reveal_type(np.dtype("float32"))  # E: numpy.dtype[numpy.float32]
+reveal_type(np.dtype("int64"))  # E: numpy.dtype[numpy.int64]
+reveal_type(np.dtype("int32"))  # E: numpy.dtype[numpy.int32]
+
+# Python types
+reveal_type(np.dtype(float))  # E: numpy.dtype[numpy.float64]
+reveal_type(np.dtype(int))  # E: numpy.dtype
+reveal_type(np.dtype(bool))  # E: numpy.dtype[numpy.bool_]
+
+# Special case for None
+reveal_type(np.dtype(None))  # E: numpy.dtype[numpy.float64]
+
+# Dtypes of dtypes
+reveal_type(np.dtype(np.dtype(np.float64)))  # E: numpy.dtype[numpy.float64*]
+
+# Void
+reveal_type(np.dtype(("U", 10)))  # E: numpy.dtype[numpy.void]