Merge pull request numpy#19275 from BvB93/dtypes

MAINT: Misc typing maintenance for `np.dtype`

numpy/__init__.pyi

@@ -5,7 +5,7 @@ import mmap
 import array as _array
 import datetime as dt
 from abc import abstractmethod
-from types import TracebackType
+from types import TracebackType, MappingProxyType
 from contextlib import ContextDecorator
 
 from numpy.core._internal import _ctypes
@@ -897,7 +897,7 @@ _DTypeScalar_co = TypeVar("_DTypeScalar_co", covariant=True, bound=generic)
 _ByteOrder = L["S", "<", ">", "=", "|", "L", "B", "N", "I"]
 
 class dtype(Generic[_DTypeScalar_co]):
-    names: Optional[Tuple[str, ...]]
+    names: None | Tuple[str, ...]
     # Overload for subclass of generic
     @overload
     def __new__(
@@ -920,7 +920,7 @@ class dtype(Generic[_DTypeScalar_co]):
     @overload
     def __new__(cls, dtype: Type[int], align: bool = ..., copy: bool = ...) -> dtype[int_]: ...
     @overload
-    def __new__(cls, dtype: Optional[Type[float]], align: bool = ..., copy: bool = ...) -> dtype[float_]: ...
+    def __new__(cls, dtype: None | Type[float], align: bool = ..., copy: bool = ...) -> dtype[float_]: ...
     @overload
     def __new__(cls, dtype: Type[complex], align: bool = ..., copy: bool = ...) -> dtype[complex_]: ...
     @overload
@@ -1051,19 +1051,23 @@ class dtype(Generic[_DTypeScalar_co]):
     @overload
     def __getitem__(self: dtype[void], key: List[str]) -> dtype[void]: ...
     @overload
-    def __getitem__(self: dtype[void], key: Union[str, int]) -> dtype[Any]: ...
+    def __getitem__(self: dtype[void], key: str | SupportsIndex) -> dtype[Any]: ...
 
-    # NOTE: In the future 1-based multiplications will also yield `void` dtypes
-    @overload
-    def __mul__(self, value: L[0]) -> None: ...  # type: ignore[misc]
+    # NOTE: In the future 1-based multiplications will also yield `flexible` dtypes
     @overload
     def __mul__(self: _DType, value: L[1]) -> _DType: ...
     @overload
-    def __mul__(self, value: int) -> dtype[void]: ...
+    def __mul__(self: _FlexDType, value: SupportsIndex) -> _FlexDType: ...
+    @overload
+    def __mul__(self, value: SupportsIndex) -> dtype[void]: ...
 
     # NOTE: `__rmul__` seems to be broken when used in combination with
-    # literals as of mypy 0.800. Set the return-type to `Any` for now.
-    def __rmul__(self, value: int) -> Any: ...
+    # literals as of mypy 0.902. Set the return-type to `dtype[Any]` for
+    # now for non-flexible dtypes.
+    @overload
+    def __rmul__(self: _FlexDType, value: SupportsIndex) -> _FlexDType: ...
+    @overload
+    def __rmul__(self, value: SupportsIndex) -> dtype[Any]: ...
 
     def __gt__(self, other: DTypeLike) -> bool: ...
     def __ge__(self, other: DTypeLike) -> bool: ...
@@ -1072,17 +1076,17 @@ class dtype(Generic[_DTypeScalar_co]):
     @property
     def alignment(self) -> int: ...
     @property
-    def base(self: _DType) -> _DType: ...
+    def base(self) -> dtype[Any]: ...
     @property
     def byteorder(self) -> str: ...
     @property
     def char(self) -> str: ...
     @property
-    def descr(self) -> List[Union[Tuple[str, str], Tuple[str, str, _Shape]]]: ...
+    def descr(self) -> List[Tuple[str, str] | Tuple[str, str, _Shape]]: ...
     @property
     def fields(
         self,
-    ) -> Optional[Mapping[str, Union[Tuple[dtype[Any], int], Tuple[dtype[Any], int, Any]]]]: ...
+    ) -> None | MappingProxyType[str, Tuple[dtype[Any], int] | Tuple[dtype[Any], int, Any]]: ...
     @property
     def flags(self) -> int: ...
     @property
@@ -1098,19 +1102,17 @@ class dtype(Generic[_DTypeScalar_co]):
     @property
     def kind(self) -> str: ...
     @property
-    def metadata(self) -> Optional[Mapping[str, Any]]: ...
+    def metadata(self) -> None | MappingProxyType[str, Any]: ...
     @property
     def name(self) -> str: ...
     @property
-    def names(self) -> Optional[Tuple[str, ...]]: ...
-    @property
     def num(self) -> int: ...
     @property
     def shape(self) -> _Shape: ...
     @property
     def ndim(self) -> int: ...
     @property
-    def subdtype(self: _DType) -> Optional[Tuple[_DType, _Shape]]: ...
+    def subdtype(self) -> None | Tuple[dtype[Any], _Shape]: ...
     def newbyteorder(self: _DType, __new_order: _ByteOrder = ...) -> _DType: ...
     # Leave str and type for end to avoid having to use `builtins.str`
     # everywhere. See https://github.com/python/mypy/issues/3775
@@ -1594,6 +1596,7 @@ class _ArrayOrScalarCommon:
 
 _DType = TypeVar("_DType", bound=dtype[Any])
 _DType_co = TypeVar("_DType_co", covariant=True, bound=dtype[Any])
+_FlexDType = TypeVar("_FlexDType", bound=dtype[flexible])
 
 # TODO: Set the `bound` to something more suitable once we
 # have proper shape support

numpy/typing/tests/data/reveal/dtype.py

@@ -1,7 +1,8 @@
 import numpy as np
 
-dtype_obj: np.dtype[np.str_]
-void_dtype_obj: np.dtype[np.void]
+dtype_U: np.dtype[np.str_]
+dtype_V: np.dtype[np.void]
+dtype_i8: np.dtype[np.int64]
 
 reveal_type(np.dtype(np.float64))  # E: numpy.dtype[{float64}]
 reveal_type(np.dtype(np.int64))  # E: numpy.dtype[{int64}]
@@ -36,22 +37,30 @@
 reveal_type(np.dtype(("U", 10)))  # E: numpy.dtype[numpy.void]
 
 # Methods and attributes
-reveal_type(dtype_obj.base)  # E: numpy.dtype[numpy.str_]
-reveal_type(dtype_obj.subdtype)  # E: Union[Tuple[numpy.dtype[numpy.str_], builtins.tuple[builtins.int]], None]
-reveal_type(dtype_obj.newbyteorder())  # E: numpy.dtype[numpy.str_]
-reveal_type(dtype_obj.type)  # E: Type[numpy.str_]
-reveal_type(dtype_obj.name)  # E: str
-reveal_type(dtype_obj.names)  # E: Union[builtins.tuple[builtins.str], None]
-
-reveal_type(dtype_obj * 0)  # E: None
-reveal_type(dtype_obj * 1)  # E: numpy.dtype[numpy.str_]
-reveal_type(dtype_obj * 2)  # E: numpy.dtype[numpy.void]
-
-reveal_type(0 * dtype_obj)  # E: Any
-reveal_type(1 * dtype_obj)  # E: Any
-reveal_type(2 * dtype_obj)  # E: Any
-
-reveal_type(void_dtype_obj["f0"])  # E: numpy.dtype[Any]
-reveal_type(void_dtype_obj[0])  # E: numpy.dtype[Any]
-reveal_type(void_dtype_obj[["f0", "f1"]])  # E: numpy.dtype[numpy.void]
-reveal_type(void_dtype_obj[["f0"]])  # E: numpy.dtype[numpy.void]
+reveal_type(dtype_U.base)  # E: numpy.dtype[Any]
+reveal_type(dtype_U.subdtype)  # E: Union[None, Tuple[numpy.dtype[Any], builtins.tuple[builtins.int]]]
+reveal_type(dtype_U.newbyteorder())  # E: numpy.dtype[numpy.str_]
+reveal_type(dtype_U.type)  # E: Type[numpy.str_]
+reveal_type(dtype_U.name)  # E: str
+reveal_type(dtype_U.names)  # E: Union[None, builtins.tuple[builtins.str]]
+
+reveal_type(dtype_U * 0)  # E: numpy.dtype[numpy.str_]
+reveal_type(dtype_U * 1)  # E: numpy.dtype[numpy.str_]
+reveal_type(dtype_U * 2)  # E: numpy.dtype[numpy.str_]
+
+reveal_type(dtype_i8 * 0)  # E: numpy.dtype[numpy.void]
+reveal_type(dtype_i8 * 1)  # E: numpy.dtype[{int64}]
+reveal_type(dtype_i8 * 2)  # E: numpy.dtype[numpy.void]
+
+reveal_type(0 * dtype_U)  # E: numpy.dtype[numpy.str_]
+reveal_type(1 * dtype_U)  # E: numpy.dtype[numpy.str_]
+reveal_type(2 * dtype_U)  # E: numpy.dtype[numpy.str_]
+
+reveal_type(0 * dtype_i8)  # E: numpy.dtype[Any]
+reveal_type(1 * dtype_i8)  # E: numpy.dtype[Any]
+reveal_type(2 * dtype_i8)  # E: numpy.dtype[Any]
+
+reveal_type(dtype_V["f0"])  # E: numpy.dtype[Any]
+reveal_type(dtype_V[0])  # E: numpy.dtype[Any]
+reveal_type(dtype_V[["f0", "f1"]])  # E: numpy.dtype[numpy.void]
+reveal_type(dtype_V[["f0"]])  # E: numpy.dtype[numpy.void]