Merge pull request numpy#28506 from jorenham/numtype/165

TYP: fix typing errors in ``numpy.lib._arrayterator_impl``

Author: charris

numpy/lib/_arrayterator_impl.pyi

@@ -1,50 +1,45 @@
+# pyright: reportIncompatibleMethodOverride=false
+
 from collections.abc import Generator
 from types import EllipsisType
-from typing import (
-    Any,
-    TypeAlias,
-    TypeVar,
-    overload,
-)
+from typing import Any, Final, TypeAlias, overload
+
+from typing_extensions import TypeVar
 
-from numpy import ndarray, dtype, generic
-from numpy._typing import DTypeLike, NDArray, _Shape as _AnyShape
+import numpy as np
 
 __all__ = ["Arrayterator"]
 
-# TODO: Rename to ``_ShapeType``
-_Shape = TypeVar("_Shape", bound=_AnyShape)
-_DType = TypeVar("_DType", bound=dtype[Any])
-_ScalarType = TypeVar("_ScalarType", bound=generic)
+_ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], covariant=True)
+_DTypeT = TypeVar("_DTypeT", bound=np.dtype[Any])
+_DTypeT_co = TypeVar("_DTypeT_co", bound=np.dtype[Any], covariant=True)
+_ScalarT = TypeVar("_ScalarT", bound=np.generic)
 
-_Index: TypeAlias = (
-    EllipsisType
-    | int
-    | slice
-    | tuple[EllipsisType | int | slice, ...]
-)
+_AnyIndex: TypeAlias = EllipsisType | int | slice | tuple[EllipsisType | int | slice, ...]
 
 # NOTE: In reality `Arrayterator` does not actually inherit from `ndarray`,
 # but its ``__getattr__` method does wrap around the former and thus has
 # access to all its methods
 
-class Arrayterator(ndarray[_Shape, _DType]):
-    var: ndarray[_Shape, _DType]  # type: ignore[assignment]
-    buf_size: None | int
-    start: list[int]
-    stop: list[int]
-    step: list[int]
+class Arrayterator(np.ndarray[_ShapeT_co, _DTypeT_co]):
+    var: np.ndarray[_ShapeT_co, _DTypeT_co]  # type: ignore[assignment]
+    buf_size: Final[int | None]
+    start: Final[list[int]]
+    stop: Final[list[int]]
+    step: Final[list[int]]
 
     @property  # type: ignore[misc]
-    def shape(self) -> tuple[int, ...]: ...
+    def shape(self) -> _ShapeT_co: ...
     @property
-    def flat(self: NDArray[_ScalarType]) -> Generator[_ScalarType, None, None]: ...
-    def __init__(
-        self, var: ndarray[_Shape, _DType], buf_size: None | int = ...
-    ) -> None: ...
-    @overload
-    def __array__(self, dtype: None = ..., copy: None | bool = ...) -> ndarray[_AnyShape, _DType]: ...
+    def flat(self: Arrayterator[Any, np.dtype[_ScalarT]]) -> Generator[_ScalarT]: ...  # type: ignore[override]
+
+    #
+    def __init__(self, /, var: np.ndarray[_ShapeT_co, _DTypeT_co], buf_size: int | None = None) -> None: ...
+    def __getitem__(self, index: _AnyIndex, /) -> Arrayterator[tuple[int, ...], _DTypeT_co]: ...  # type: ignore[override]
+    def __iter__(self) -> Generator[np.ndarray[tuple[int, ...], _DTypeT_co]]: ...
+
+    #
+    @overload  # type: ignore[override]
+    def __array__(self, /, dtype: None = None, copy: bool | None = None) -> np.ndarray[_ShapeT_co, _DTypeT_co]: ...
     @overload
-    def __array__(self, dtype: DTypeLike, copy: None | bool = ...) -> NDArray[Any]: ...
-    def __getitem__(self, index: _Index) -> Arrayterator[_AnyShape, _DType]: ...
-    def __iter__(self) -> Generator[ndarray[_AnyShape, _DType], None, None]: ...
+    def __array__(self, /, dtype: _DTypeT, copy: bool | None = None) -> np.ndarray[_ShapeT_co, _DTypeT]: ...