Merge pull request numpy#26928 from jorenham/covariant-flatiter

TYP: Covariant `numpy.flatiter` type parameter

Author: rgommers

numpy/__init__.pyi

@@ -638,6 +638,7 @@ test: PytestTester
 def show_config() -> None: ...
 
 _NdArraySubClass = TypeVar("_NdArraySubClass", bound=NDArray[Any])
+_NdArraySubClass_co = TypeVar("_NdArraySubClass_co", bound=NDArray[Any], covariant=True)
 _DTypeScalar_co = TypeVar("_DTypeScalar_co", covariant=True, bound=generic)
 _ByteOrder: TypeAlias = L["S", "<", ">", "=", "|", "L", "B", "N", "I", "little", "big", "native"]
 
@@ -904,15 +905,15 @@ _FlatIterSelf = TypeVar("_FlatIterSelf", bound=flatiter[Any])
 _FlatShapeType = TypeVar("_FlatShapeType", bound=tuple[int])
 
 @final
-class flatiter(Generic[_NdArraySubClass]):
+class flatiter(Generic[_NdArraySubClass_co]):
     __hash__: ClassVar[None]
     @property
-    def base(self) -> _NdArraySubClass: ...
+    def base(self) -> _NdArraySubClass_co: ...
     @property
     def coords(self) -> _Shape: ...
     @property
     def index(self) -> int: ...
-    def copy(self) -> _NdArraySubClass: ...
+    def copy(self) -> _NdArraySubClass_co: ...
     def __iter__(self: _FlatIterSelf) -> _FlatIterSelf: ...
     def __next__(self: flatiter[NDArray[_ScalarType]]) -> _ScalarType: ...
     def __len__(self) -> int: ...
@@ -925,7 +926,7 @@ class flatiter(Generic[_NdArraySubClass]):
     def __getitem__(
         self,
         key: _ArrayLikeInt | slice | ellipsis | tuple[_ArrayLikeInt | slice | ellipsis],
-    ) -> _NdArraySubClass: ...
+    ) -> _NdArraySubClass_co: ...
     # TODO: `__setitem__` operates via `unsafe` casting rules, and can
     # thus accept any type accepted by the relevant underlying `np.generic`
     # constructor.