1  from collections.abc import Generator
 2  from typing import (
 3      Any,
 4      TypeVar,
 5      Union,
 6      overload,
 7  )
 8  
 9  from numpy import ndarray, dtype, generic
10  from numpy._typing import DTypeLike
11  
12  # TODO: Set a shape bound once we've got proper shape support
13  _Shape = TypeVar("_Shape", bound=Any)
14  _DType = TypeVar("_DType", bound=dtype[Any])
15  _ScalarType = TypeVar("_ScalarType", bound=generic)
16  
17  _Index = Union[
18      Union[ellipsis, int, slice],
19      tuple[Union[ellipsis, int, slice], ...],
20  ]
21  
22  __all__: list[str]
23  
24  # NOTE: In reality `Arrayterator` does not actually inherit from `ndarray`,
25  # but its ``__getattr__` method does wrap around the former and thus has
26  # access to all its methods
27  
28  class Arrayterator(ndarray[_Shape, _DType]):
29      var: ndarray[_Shape, _DType]  # type: ignore[assignment]
30      buf_size: None | int
31      start: list[int]
32      stop: list[int]
33      step: list[int]
34  
35      @property  # type: ignore[misc]
36      def shape(self) -> tuple[int, ...]: ...
37      @property
38      def flat(  # type: ignore[override]
39          self: ndarray[Any, dtype[_ScalarType]]
40      ) -> Generator[_ScalarType, None, None]: ...
41      def __init__(
42          self, var: ndarray[_Shape, _DType], buf_size: None | int = ...
43      ) -> None: ...
44      @overload
45      def __array__(self, dtype: None = ...) -> ndarray[Any, _DType]: ...
46      @overload
47      def __array__(self, dtype: DTypeLike) -> ndarray[Any, dtype[Any]]: ...
48      def __getitem__(self, index: _Index) -> Arrayterator[Any, _DType]: ...
49      def __iter__(self) -> Generator[ndarray[Any, _DType], None, None]: ...