codekingpro/portable-devtools

1# pyright: reportIncompatibleMethodOverride=false
2
3from collections.abc import Generator
4from types import EllipsisType
5from typing import Any, Final, TypeAlias, overload
6from typing_extensions import TypeVar
7
8import numpy as np
9from numpy._typing import _AnyShape, _Shape
10
11__all__ = ["Arrayterator"]
12
13_ShapeT_co = TypeVar("_ShapeT_co", bound=_Shape, default=_AnyShape, covariant=True)
14_DTypeT = TypeVar("_DTypeT", bound=np.dtype)
15_DTypeT_co = TypeVar("_DTypeT_co", bound=np.dtype, default=np.dtype, covariant=True)
16_ScalarT = TypeVar("_ScalarT", bound=np.generic)
17
18_AnyIndex: TypeAlias = EllipsisType | int | slice | tuple[EllipsisType | int | slice, ...]
19
20# NOTE: In reality `Arrayterator` does not actually inherit from `ndarray`,
21# but its ``__getattr__` method does wrap around the former and thus has
22# access to all its methods
23
24class Arrayterator(np.ndarray[_ShapeT_co, _DTypeT_co]):
25    var: np.ndarray[_ShapeT_co, _DTypeT_co]  # type: ignore[assignment]
26    buf_size: Final[int | None]
27    start: Final[list[int]]
28    stop: Final[list[int]]
29    step: Final[list[int]]
30
31    @property  # type: ignore[misc]
32    def shape(self) -> _ShapeT_co: ...
33    @property
34    def flat(self: Arrayterator[Any, np.dtype[_ScalarT]]) -> Generator[_ScalarT]: ...  # type: ignore[override]
35
36    #
37    def __init__(self, /, var: np.ndarray[_ShapeT_co, _DTypeT_co], buf_size: int | None = None) -> None: ...
38    def __getitem__(self, index: _AnyIndex, /) -> Arrayterator[_AnyShape, _DTypeT_co]: ...  # type: ignore[override]
39    def __iter__(self) -> Generator[np.ndarray[_AnyShape, _DTypeT_co]]: ...
40
41    #
42    @overload  # type: ignore[override]
43    def __array__(self, /, dtype: None = None, copy: bool | None = None) -> np.ndarray[_ShapeT_co, _DTypeT_co]: ...
44    @overload
45    def __array__(self, /, dtype: _DTypeT, copy: bool | None = None) -> np.ndarray[_ShapeT_co, _DTypeT]: ...
46