arrayterator.pyi 1.5 KB

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