Merge pull request #17673 from BvB93/generic-specific

MAINT: Refine a number of `np.generic` annotations

4 files changed, 96 insertions, 47 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index de7efdcaa..579dd03d6 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -999,25 +999,13 @@ class _ArrayOrScalarCommon:
     @property
     def T(self: _ArraySelf) -> _ArraySelf: ...
     @property
-    def base(self) -> Optional[ndarray]: ...
-    @property
-    def dtype(self) -> _DType: ...
-    @property
     def data(self) -> memoryview: ...
     @property
     def flags(self) -> _flagsobj: ...
     @property
-    def size(self) -> int: ...
-    @property
     def itemsize(self) -> int: ...
     @property
     def nbytes(self) -> int: ...
-    @property
-    def ndim(self) -> int: ...
-    @property
-    def shape(self) -> _Shape: ...
-    @property
-    def strides(self) -> _Shape: ...
     def __array__(self, __dtype: DTypeLike = ...) -> ndarray: ...
     def __bool__(self) -> bool: ...
     def __bytes__(self) -> bytes: ...
@@ -1039,45 +1027,22 @@ class _ArrayOrScalarCommon:
         subok: bool = ...,
         copy: bool = ...,
     ) -> _ArraySelf: ...
-    def byteswap(self: _ArraySelf, inplace: bool = ...) -> _ArraySelf: ...
     def copy(self: _ArraySelf, order: _OrderKACF = ...) -> _ArraySelf: ...
     def dump(self, file: str) -> None: ...
     def dumps(self) -> bytes: ...
-    def fill(self, value: Any) -> None: ...
-    @property
-    def flat(self: _ArraySelf) -> flatiter[_ArraySelf]: ...
-    def flatten(self: _ArraySelf, order: _OrderKACF = ...) -> _ArraySelf: ...
+    def flatten(self, order: _OrderKACF = ...) -> ndarray: ...
     def getfield(
         self: _ArraySelf, dtype: DTypeLike, offset: int = ...
     ) -> _ArraySelf: ...
-    @overload
-    def item(self, *args: int) -> Any: ...
-    @overload
-    def item(self, args: Tuple[int, ...]) -> Any: ...
-    @overload
-    def itemset(self, __value: Any) -> None: ...
-    @overload
-    def itemset(self, __item: _ShapeLike, __value: Any) -> None: ...
-    def ravel(self: _ArraySelf, order: _OrderKACF = ...) -> _ArraySelf: ...
+    def ravel(self, order: _OrderKACF = ...) -> ndarray: ...
     @overload
     def reshape(
-        self: _ArraySelf, shape: Sequence[int], *, order: _OrderACF = ...
-    ) -> _ArraySelf: ...
+        self, __shape: Sequence[int], *, order: _OrderACF = ...
+    ) -> ndarray: ...
     @overload
     def reshape(
-        self: _ArraySelf, *shape: int, order: _OrderACF = ...
-    ) -> _ArraySelf: ...
-    @overload
-    def resize(self, new_shape: Sequence[int], *, refcheck: bool = ...) -> None: ...
-    @overload
-    def resize(self, *new_shape: int, refcheck: bool = ...) -> None: ...
-    def setflags(
-        self, write: bool = ..., align: bool = ..., uic: bool = ...
-    ) -> None: ...
-    def squeeze(
-        self: _ArraySelf, axis: Union[int, Tuple[int, ...]] = ...
-    ) -> _ArraySelf: ...
-    def swapaxes(self: _ArraySelf, axis1: int, axis2: int) -> _ArraySelf: ...
+        self, *shape: int, order: _OrderACF = ...
+    ) -> ndarray: ...
     def tobytes(self, order: _OrderKACF = ...) -> bytes: ...
     # NOTE: `tostring()` is deprecated and therefore excluded
     # def tostring(self, order=...): ...
@@ -1087,10 +1052,6 @@ class _ArrayOrScalarCommon:
     # generics and 0d arrays return builtin scalars
     def tolist(self) -> Any: ...
     @overload
-    def transpose(self: _ArraySelf, axes: Sequence[int]) -> _ArraySelf: ...
-    @overload
-    def transpose(self: _ArraySelf, *axes: int) -> _ArraySelf: ...
-    @overload
     def view(self, type: Type[_NdArraySubClass]) -> _NdArraySubClass: ...
     @overload
     def view(self: _ArraySelf, dtype: DTypeLike = ...) -> _ArraySelf: ...
@@ -1485,6 +1446,12 @@ _Casting = Literal["no", "equiv", "safe", "same_kind", "unsafe"]
 
 class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container):
     @property
+    def base(self) -> Optional[ndarray]: ...
+    @property
+    def ndim(self) -> int: ...
+    @property
+    def size(self) -> int: ...
+    @property
     def real(self: _ArraySelf) -> _ArraySelf: ...
     @real.setter
     def real(self, value: ArrayLike) -> None: ...
@@ -1513,6 +1480,33 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container):
     def strides(self) -> _Shape: ...
     @strides.setter
     def strides(self, value: _ShapeLike): ...
+    def byteswap(self: _ArraySelf, inplace: bool = ...) -> _ArraySelf: ...
+    def fill(self, value: Any) -> None: ...
+    @property
+    def flat(self: _ArraySelf) -> flatiter[_ArraySelf]: ...
+    @overload
+    def item(self, *args: int) -> Any: ...
+    @overload
+    def item(self, __args: Tuple[int, ...]) -> Any: ...
+    @overload
+    def itemset(self, __value: Any) -> None: ...
+    @overload
+    def itemset(self, __item: _ShapeLike, __value: Any) -> None: ...
+    @overload
+    def resize(self, __new_shape: Sequence[int], *, refcheck: bool = ...) -> None: ...
+    @overload
+    def resize(self, *new_shape: int, refcheck: bool = ...) -> None: ...
+    def setflags(
+        self, write: bool = ..., align: bool = ..., uic: bool = ...
+    ) -> None: ...
+    def squeeze(
+        self: _ArraySelf, axis: Union[int, Tuple[int, ...]] = ...
+    ) -> _ArraySelf: ...
+    def swapaxes(self: _ArraySelf, axis1: int, axis2: int) -> _ArraySelf: ...
+    @overload
+    def transpose(self: _ArraySelf, __axes: Sequence[int]) -> _ArraySelf: ...
+    @overload
+    def transpose(self: _ArraySelf, *axes: int) -> _ArraySelf: ...
     def argpartition(
         self,
         kth: _ArrayLikeIntOrBool,
@@ -1644,6 +1638,7 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container):
 
 # See https://github.com/numpy/numpy-stubs/pull/80 for more details.
 
+_ScalarType = TypeVar("_ScalarType", bound=generic)
 _NBit_co = TypeVar("_NBit_co", covariant=True, bound=NBitBase)
 _NBit_co2 = TypeVar("_NBit_co2", covariant=True, bound=NBitBase)
 
@@ -1652,6 +1647,27 @@ class generic(_ArrayOrScalarCommon):
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     @property
     def base(self) -> None: ...
+    @property
+    def dtype(self: _ScalarType) -> _DType[_ScalarType]: ...
+    @property
+    def ndim(self) -> Literal[0]: ...
+    @property
+    def size(self) -> Literal[1]: ...
+    @property
+    def shape(self) -> Tuple[()]: ...
+    @property
+    def strides(self) -> Tuple[()]: ...
+    def byteswap(self: _ScalarType, inplace: Literal[False] = ...) -> _ScalarType: ...
+    @property
+    def flat(self) -> flatiter[ndarray]: ...
+    def item(
+        self: _ScalarType,
+        __args: Union[Literal[0], Tuple[()], Tuple[Literal[0]]] = ...,
+    ) -> Any: ...
+    def squeeze(
+        self: _ScalarType, axis: Union[Literal[0], Tuple[()]] = ...
+    ) -> _ScalarType: ...
+    def transpose(self: _ScalarType, __axes: Tuple[()] = ...) -> _ScalarType: ...
 
 class number(generic, Generic[_NBit_co]):  # type: ignore
     @property
diff --git a/numpy/typing/tests/data/fail/scalars.py b/numpy/typing/tests/data/fail/scalars.py
index 5fffa89b1..f09740875 100644
--- a/numpy/typing/tests/data/fail/scalars.py
+++ b/numpy/typing/tests/data/fail/scalars.py
@@ -1,5 +1,7 @@
 import numpy as np
 
+f8: np.float64
+
 # Construction
 
 np.float32(3j)  # E: incompatible type
@@ -66,3 +68,9 @@ np.bytes_(b"hello", encoding='utf-8')  # E: No overload variant
 np.str_("hello", encoding='utf-8')  # E: No overload variant
 
 complex(np.bytes_("1"))  # E: No overload variant
+
+f8.item(1)  # E: incompatible type
+f8.item((0, 1))  # E: incompatible type
+f8.squeeze(axis=1)  # E: incompatible type
+f8.squeeze(axis=(0, 1))  # E: incompatible type
+f8.transpose(1)  # E: incompatible type
diff --git a/numpy/typing/tests/data/pass/scalars.py b/numpy/typing/tests/data/pass/scalars.py
index 04e974983..b7f7880e4 100644
--- a/numpy/typing/tests/data/pass/scalars.py
+++ b/numpy/typing/tests/data/pass/scalars.py
@@ -147,3 +147,19 @@ complex(b_)
 complex(td)
 complex(U)
 complex(AR)
+
+
+# Misc
+c16.dtype
+c16.real
+c16.imag
+c16.real.real
+c16.real.imag
+c16.ndim
+c16.size
+c16.itemsize
+c16.shape
+c16.strides
+c16.squeeze()
+c16.byteswap()
+c16.transpose()
diff --git a/numpy/typing/tests/data/reveal/scalars.py b/numpy/typing/tests/data/reveal/scalars.py
index 72a3e4d2e..e887e302d 100644
--- a/numpy/typing/tests/data/reveal/scalars.py
+++ b/numpy/typing/tests/data/reveal/scalars.py
@@ -9,8 +9,17 @@ reveal_type(x.real.real)  # E: numpy.floating[numpy.typing._32Bit]
 reveal_type(x.real.imag)  # E: numpy.floating[numpy.typing._32Bit]
 
 reveal_type(x.itemsize)  # E: int
-reveal_type(x.shape)  # E: tuple[builtins.int]
-reveal_type(x.strides)  # E: tuple[builtins.int]
+reveal_type(x.shape)  # E: Tuple[]
+reveal_type(x.strides)  # E: Tuple[]
+
+reveal_type(x.ndim)  # E: Literal[0]
+reveal_type(x.size)  # E: Literal[1]
+
+reveal_type(x.squeeze())  # E: numpy.complexfloating[numpy.typing._32Bit, numpy.typing._32Bit]
+reveal_type(x.byteswap())  # E: numpy.complexfloating[numpy.typing._32Bit, numpy.typing._32Bit]
+reveal_type(x.transpose())  # E: numpy.complexfloating[numpy.typing._32Bit, numpy.typing._32Bit]
+
+reveal_type(x.dtype)  # E: numpy.dtype[numpy.complexfloating[numpy.typing._32Bit, numpy.typing._32Bit]]
 
 reveal_type(np.complex64().real)  # E: numpy.floating[numpy.typing._32Bit]
 reveal_type(np.complex128().imag)  # E: numpy.floating[numpy.typing._64Bit]