TYP: Transparent __array__ shape-type (numpy#26927)

This changes the  ndarray.__array__ and flatiter.__array__ methods to return a ndarray with the same shape type.

Due to technical limitations, flatiter will it only return the shape type of its underlying ndarray if it's 1-d (like tuple[int]).

Author: jorenham

numpy/__init__.pyi

@@ -901,6 +901,7 @@ _ArrayLikeInt: TypeAlias = (
 )
 
 _FlatIterSelf = TypeVar("_FlatIterSelf", bound=flatiter[Any])
+_FlatShapeType = TypeVar("_FlatShapeType", bound=tuple[int])
 
 @final
 class flatiter(Generic[_NdArraySubClass]):
@@ -935,6 +936,10 @@ class flatiter(Generic[_NdArraySubClass]):
         value: Any,
     ) -> None: ...
     @overload
+    def __array__(self: flatiter[ndarray[_FlatShapeType, _DType]], dtype: None = ..., /) -> ndarray[_FlatShapeType, _DType]: ...
+    @overload
+    def __array__(self: flatiter[ndarray[_FlatShapeType, Any]], dtype: _DType, /) -> ndarray[_FlatShapeType, _DType]: ...
+    @overload
     def __array__(self: flatiter[ndarray[Any, _DType]], dtype: None = ..., /) -> ndarray[Any, _DType]: ...
     @overload
     def __array__(self, dtype: _DType, /) -> ndarray[Any, _DType]: ...
@@ -1469,11 +1474,11 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
     @overload
     def __array__(
         self, dtype: None = ..., /, *, copy: None | bool = ...
-    ) -> ndarray[Any, _DType_co]: ...
+    ) -> ndarray[_ShapeType, _DType_co]: ...
     @overload
     def __array__(
         self, dtype: _DType, /, *, copy: None | bool = ...
-    ) -> ndarray[Any, _DType]: ...
+    ) -> ndarray[_ShapeType, _DType]: ...
 
     def __array_ufunc__(
         self,
@@ -1704,11 +1709,13 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
         axis: None | SupportsIndex = ...,
     ) -> ndarray[Any, _DType_co]: ...
 
+    # TODO: use `tuple[int]` as shape type once covariant (#26081)
     def flatten(
         self,
         order: _OrderKACF = ...,
     ) -> ndarray[Any, _DType_co]: ...
 
+    # TODO: use `tuple[int]` as shape type once covariant (#26081)
     def ravel(
         self,
         order: _OrderKACF = ...,
@@ -2613,6 +2620,7 @@ _NBit2 = TypeVar("_NBit2", bound=NBitBase)
 class generic(_ArrayOrScalarCommon):
     @abstractmethod
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+    # TODO: use `tuple[()]` as shape type once covariant (#26081)
     @overload
     def __array__(self: _ScalarType, dtype: None = ..., /) -> NDArray[_ScalarType]: ...
     @overload
@@ -3740,6 +3748,7 @@ class poly1d:
 
     __hash__: ClassVar[None]  # type: ignore
 
+    # TODO: use `tuple[int]` as shape type once covariant (#26081)
     @overload
     def __array__(self, t: None = ..., copy: None | bool = ...) -> NDArray[Any]: ...
     @overload

numpy/typing/tests/data/reveal/flatiter.pyi

@@ -1,5 +1,5 @@
 import sys
-from typing import Any
+from typing import Any, Literal, TypeAlias
 
 import numpy as np
 import numpy.typing as npt
@@ -10,6 +10,10 @@ else:
     from typing_extensions import assert_type
 
 a: np.flatiter[npt.NDArray[np.str_]]
+a_1d: np.flatiter[np.ndarray[tuple[int], np.dtype[np.bytes_]]]
+
+Size: TypeAlias = Literal[42]
+a_1d_fixed: np.flatiter[np.ndarray[tuple[Size], np.dtype[np.object_]]]
 
 assert_type(a.base, npt.NDArray[np.str_])
 assert_type(a.copy(), npt.NDArray[np.str_])
@@ -23,8 +27,26 @@ assert_type(a[...], npt.NDArray[np.str_])
 assert_type(a[:], npt.NDArray[np.str_])
 assert_type(a[(...,)], npt.NDArray[np.str_])
 assert_type(a[(0,)], np.str_)
+
 assert_type(a.__array__(), npt.NDArray[np.str_])
 assert_type(a.__array__(np.dtype(np.float64)), npt.NDArray[np.float64])
+assert_type(
+    a_1d.__array__(),
+    np.ndarray[tuple[int], np.dtype[np.bytes_]],
+)
+assert_type(
+    a_1d.__array__(np.dtype(np.float64)),
+    np.ndarray[tuple[int], np.dtype[np.float64]],
+)
+assert_type(
+    a_1d_fixed.__array__(),
+    np.ndarray[tuple[Size], np.dtype[np.object_]],
+)
+assert_type(
+    a_1d_fixed.__array__(np.dtype(np.float64)),
+    np.ndarray[tuple[Size], np.dtype[np.float64]],
+)
+
 a[0] = "a"
 a[:5] = "a"
 a[...] = "a"