_user_array_impl.pyi 9.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220
  1. from types import EllipsisType
  2. from typing import Any, Generic, SupportsIndex, TypeAlias, TypeVar, overload
  3. from _typeshed import Incomplete
  4. from typing_extensions import Self, deprecated, override
  5. import numpy as np
  6. import numpy.typing as npt
  7. from numpy._typing import _ArrayLike, _ArrayLikeBool_co, _ArrayLikeInt_co, _DTypeLike
  8. ###
  9. _ScalarT = TypeVar("_ScalarT", bound=np.generic)
  10. _ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
  11. _ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], default=Any, covariant=True)
  12. _DTypeT = TypeVar("_DTypeT", bound=np.dtype[Any])
  13. _DTypeT_co = TypeVar("_DTypeT_co", bound=np.dtype[Any], default=np.dtype[Any], covariant=True)
  14. _BoolArrayT = TypeVar("_BoolArrayT", bound=container[Any, np.dtype[np.bool]])
  15. _IntegralArrayT = TypeVar("_IntegralArrayT", bound=container[Any, np.dtype[np.bool | np.integer | np.object_]])
  16. _RealContainerT = TypeVar(
  17. "_RealContainerT",
  18. bound=container[Any, np.dtype[np.bool | np.integer | np.floating | np.timedelta64 | np.object_]],
  19. )
  20. _NumericContainerT = TypeVar("_NumericContainerT", bound=container[Any, np.dtype[np.number | np.timedelta64 | np.object_]])
  21. _ArrayInt_co: TypeAlias = npt.NDArray[np.integer | np.bool]
  22. _ToIndexSlice: TypeAlias = slice | EllipsisType | _ArrayInt_co | None
  23. _ToIndexSlices: TypeAlias = _ToIndexSlice | tuple[_ToIndexSlice, ...]
  24. _ToIndex: TypeAlias = SupportsIndex | _ToIndexSlice
  25. _ToIndices: TypeAlias = _ToIndex | tuple[_ToIndex, ...]
  26. ###
  27. class container(Generic[_ShapeT_co, _DTypeT_co]):
  28. array: np.ndarray[_ShapeT_co, _DTypeT_co]
  29. @overload
  30. def __init__(
  31. self,
  32. /,
  33. data: container[_ShapeT_co, _DTypeT_co] | np.ndarray[_ShapeT_co, _DTypeT_co],
  34. dtype: None = None,
  35. copy: bool = True,
  36. ) -> None: ...
  37. @overload
  38. def __init__(
  39. self: container[Any, np.dtype[_ScalarT]],
  40. /,
  41. data: _ArrayLike[_ScalarT],
  42. dtype: None = None,
  43. copy: bool = True,
  44. ) -> None: ...
  45. @overload
  46. def __init__(
  47. self: container[Any, np.dtype[_ScalarT]],
  48. /,
  49. data: npt.ArrayLike,
  50. dtype: _DTypeLike[_ScalarT],
  51. copy: bool = True,
  52. ) -> None: ...
  53. @overload
  54. def __init__(self, /, data: npt.ArrayLike, dtype: npt.DTypeLike | None = None, copy: bool = True) -> None: ...
  55. #
  56. def __complex__(self, /) -> complex: ...
  57. def __float__(self, /) -> float: ...
  58. def __int__(self, /) -> int: ...
  59. def __hex__(self, /) -> str: ...
  60. def __oct__(self, /) -> str: ...
  61. #
  62. @override
  63. def __eq__(self, other: object, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
  64. @override
  65. def __ne__(self, other: object, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
  66. #
  67. def __lt__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
  68. def __le__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
  69. def __gt__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
  70. def __ge__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
  71. #
  72. def __len__(self, /) -> int: ...
  73. # keep in sync with np.ndarray
  74. @overload
  75. def __getitem__(self, key: _ArrayInt_co | tuple[_ArrayInt_co, ...], /) -> container[_ShapeT_co, _DTypeT_co]: ...
  76. @overload
  77. def __getitem__(self, key: _ToIndexSlices, /) -> container[Any, _DTypeT_co]: ...
  78. @overload
  79. def __getitem__(self, key: _ToIndices, /) -> Any: ...
  80. @overload
  81. def __getitem__(self: container[Any, np.dtype[np.void]], key: list[str], /) -> container[_ShapeT_co, np.dtype[np.void]]: ...
  82. @overload
  83. def __getitem__(self: container[Any, np.dtype[np.void]], key: str, /) -> container[_ShapeT_co, np.dtype[Any]]: ...
  84. # keep in sync with np.ndarray
  85. @overload
  86. def __setitem__(self, index: _ToIndices, value: object, /) -> None: ...
  87. @overload
  88. def __setitem__(self: container[Any, np.dtype[np.void]], key: str | list[str], value: object, /) -> None: ...
  89. # keep in sync with np.ndarray
  90. @overload
  91. def __abs__(self: container[_ShapeT, np.dtype[np.complex64]], /) -> container[_ShapeT, np.dtype[np.float32]]: ... # type: ignore[overload-overlap]
  92. @overload
  93. def __abs__(self: container[_ShapeT, np.dtype[np.complex128]], /) -> container[_ShapeT, np.dtype[np.float64]]: ...
  94. @overload
  95. def __abs__(self: container[_ShapeT, np.dtype[np.complex192]], /) -> container[_ShapeT, np.dtype[np.float96]]: ...
  96. @overload
  97. def __abs__(self: container[_ShapeT, np.dtype[np.complex256]], /) -> container[_ShapeT, np.dtype[np.float128]]: ...
  98. @overload
  99. def __abs__(self: _RealContainerT, /) -> _RealContainerT: ...
  100. #
  101. def __neg__(self: _NumericContainerT, /) -> _NumericContainerT: ... # noqa: PYI019
  102. def __pos__(self: _NumericContainerT, /) -> _NumericContainerT: ... # noqa: PYI019
  103. def __invert__(self: _IntegralArrayT, /) -> _IntegralArrayT: ... # noqa: PYI019
  104. # TODO(jorenham): complete these binary ops
  105. #
  106. def __add__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  107. def __radd__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  108. def __iadd__(self, other: npt.ArrayLike, /) -> Self: ...
  109. #
  110. def __sub__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  111. def __rsub__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  112. def __isub__(self, other: npt.ArrayLike, /) -> Self: ...
  113. #
  114. def __mul__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  115. def __rmul__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  116. def __imul__(self, other: npt.ArrayLike, /) -> Self: ...
  117. #
  118. def __div__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  119. def __rdiv__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  120. def __idiv__(self, other: npt.ArrayLike, /) -> Self: ...
  121. #
  122. def __mod__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  123. def __rmod__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  124. def __imod__(self, other: npt.ArrayLike, /) -> Self: ...
  125. #
  126. def __divmod__(self, other: npt.ArrayLike, /) -> tuple[Incomplete, Incomplete]: ...
  127. def __rdivmod__(self, other: npt.ArrayLike, /) -> tuple[Incomplete, Incomplete]: ...
  128. #
  129. def __pow__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  130. def __rpow__(self, other: npt.ArrayLike, /) -> Incomplete: ...
  131. def __ipow__(self, other: npt.ArrayLike, /) -> Self: ...
  132. #
  133. def __lshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ...
  134. def __rlshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ...
  135. def __ilshift__(self, other: _ArrayLikeInt_co, /) -> Self: ...
  136. #
  137. def __rshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ...
  138. def __rrshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ...
  139. def __irshift__(self, other: _ArrayLikeInt_co, /) -> Self: ...
  140. #
  141. @overload
  142. def __and__(self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /) -> container[Any, np.dtype[np.bool]]: ...
  143. @overload
  144. def __and__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ...
  145. __rand__ = __and__
  146. @overload
  147. def __iand__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ...
  148. @overload
  149. def __iand__(self, other: _ArrayLikeInt_co, /) -> Self: ...
  150. #
  151. @overload
  152. def __xor__(self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /) -> container[Any, np.dtype[np.bool]]: ...
  153. @overload
  154. def __xor__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ...
  155. __rxor__ = __xor__
  156. @overload
  157. def __ixor__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ...
  158. @overload
  159. def __ixor__(self, other: _ArrayLikeInt_co, /) -> Self: ...
  160. #
  161. @overload
  162. def __or__(self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /) -> container[Any, np.dtype[np.bool]]: ...
  163. @overload
  164. def __or__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ...
  165. __ror__ = __or__
  166. @overload
  167. def __ior__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ...
  168. @overload
  169. def __ior__(self, other: _ArrayLikeInt_co, /) -> Self: ...
  170. #
  171. @overload
  172. def __array__(self, /, t: None = None) -> np.ndarray[_ShapeT_co, _DTypeT_co]: ...
  173. @overload
  174. def __array__(self, /, t: _DTypeT) -> np.ndarray[_ShapeT_co, _DTypeT]: ...
  175. #
  176. @overload
  177. def __array_wrap__(self, arg0: npt.ArrayLike, /) -> container[_ShapeT_co, _DTypeT_co]: ...
  178. @overload
  179. def __array_wrap__(self, a: np.ndarray[_ShapeT, _DTypeT], c: Any = ..., s: Any = ..., /) -> container[_ShapeT, _DTypeT]: ...
  180. #
  181. def copy(self, /) -> Self: ...
  182. @deprecated("tostring() is deprecated. Use tobytes() instead.")
  183. def tostring(self, /) -> bytes: ...
  184. def tobytes(self, /) -> bytes: ...
  185. def byteswap(self, /) -> Self: ...
  186. def astype(self, /, typecode: _DTypeLike[_ScalarT]) -> container[_ShapeT_co, np.dtype[_ScalarT]]: ...