_arrayterator_impl.pyi 1.8 KB

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