index_tricks.pyi 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. import sys
  2. from typing import (
  3. Any,
  4. Tuple,
  5. TypeVar,
  6. Generic,
  7. overload,
  8. List,
  9. Union,
  10. Sequence,
  11. )
  12. from numpy import (
  13. # Circumvent a naming conflict with `AxisConcatenator.matrix`
  14. matrix as _Matrix,
  15. ndenumerate as ndenumerate,
  16. ndindex as ndindex,
  17. ndarray,
  18. dtype,
  19. integer,
  20. str_,
  21. bytes_,
  22. bool_,
  23. int_,
  24. float_,
  25. complex_,
  26. intp,
  27. _OrderCF,
  28. _ModeKind,
  29. )
  30. from numpy.typing import (
  31. # Arrays
  32. ArrayLike,
  33. _NestedSequence,
  34. _RecursiveSequence,
  35. NDArray,
  36. _ArrayLikeInt,
  37. # DTypes
  38. DTypeLike,
  39. _SupportsDType,
  40. # Shapes
  41. _ShapeLike,
  42. )
  43. if sys.version_info >= (3, 8):
  44. from typing import Literal, SupportsIndex
  45. else:
  46. from typing_extensions import Literal, SupportsIndex
  47. _T = TypeVar("_T")
  48. _DType = TypeVar("_DType", bound=dtype[Any])
  49. _BoolType = TypeVar("_BoolType", Literal[True], Literal[False])
  50. _TupType = TypeVar("_TupType", bound=Tuple[Any, ...])
  51. _ArrayType = TypeVar("_ArrayType", bound=ndarray[Any, Any])
  52. __all__: List[str]
  53. @overload
  54. def unravel_index( # type: ignore[misc]
  55. indices: Union[int, integer[Any]],
  56. shape: _ShapeLike,
  57. order: _OrderCF = ...
  58. ) -> Tuple[intp, ...]: ...
  59. @overload
  60. def unravel_index(
  61. indices: _ArrayLikeInt,
  62. shape: _ShapeLike,
  63. order: _OrderCF = ...
  64. ) -> Tuple[NDArray[intp], ...]: ...
  65. @overload
  66. def ravel_multi_index( # type: ignore[misc]
  67. multi_index: Sequence[Union[int, integer[Any]]],
  68. dims: _ShapeLike,
  69. mode: Union[_ModeKind, Tuple[_ModeKind, ...]] = ...,
  70. order: _OrderCF = ...
  71. ) -> intp: ...
  72. @overload
  73. def ravel_multi_index(
  74. multi_index: Sequence[_ArrayLikeInt],
  75. dims: _ShapeLike,
  76. mode: Union[_ModeKind, Tuple[_ModeKind, ...]] = ...,
  77. order: _OrderCF = ...
  78. ) -> NDArray[intp]: ...
  79. @overload
  80. def ix_(*args: _NestedSequence[_SupportsDType[_DType]]) -> Tuple[ndarray[Any, _DType], ...]: ...
  81. @overload
  82. def ix_(*args: _NestedSequence[str]) -> Tuple[NDArray[str_], ...]: ...
  83. @overload
  84. def ix_(*args: _NestedSequence[bytes]) -> Tuple[NDArray[bytes_], ...]: ...
  85. @overload
  86. def ix_(*args: _NestedSequence[bool]) -> Tuple[NDArray[bool_], ...]: ...
  87. @overload
  88. def ix_(*args: _NestedSequence[int]) -> Tuple[NDArray[int_], ...]: ...
  89. @overload
  90. def ix_(*args: _NestedSequence[float]) -> Tuple[NDArray[float_], ...]: ...
  91. @overload
  92. def ix_(*args: _NestedSequence[complex]) -> Tuple[NDArray[complex_], ...]: ...
  93. @overload
  94. def ix_(*args: _RecursiveSequence) -> Tuple[NDArray[Any], ...]: ...
  95. class nd_grid(Generic[_BoolType]):
  96. sparse: _BoolType
  97. def __init__(self, sparse: _BoolType = ...) -> None: ...
  98. @overload
  99. def __getitem__(
  100. self: nd_grid[Literal[False]],
  101. key: Union[slice, Sequence[slice]],
  102. ) -> NDArray[Any]: ...
  103. @overload
  104. def __getitem__(
  105. self: nd_grid[Literal[True]],
  106. key: Union[slice, Sequence[slice]],
  107. ) -> List[NDArray[Any]]: ...
  108. class MGridClass(nd_grid[Literal[False]]):
  109. def __init__(self) -> None: ...
  110. mgrid: MGridClass
  111. class OGridClass(nd_grid[Literal[True]]):
  112. def __init__(self) -> None: ...
  113. ogrid: OGridClass
  114. class AxisConcatenator:
  115. axis: int
  116. matrix: bool
  117. ndmin: int
  118. trans1d: int
  119. def __init__(
  120. self,
  121. axis: int = ...,
  122. matrix: bool = ...,
  123. ndmin: int = ...,
  124. trans1d: int = ...,
  125. ) -> None: ...
  126. @staticmethod
  127. @overload
  128. def concatenate( # type: ignore[misc]
  129. *a: ArrayLike, axis: SupportsIndex = ..., out: None = ...
  130. ) -> NDArray[Any]: ...
  131. @staticmethod
  132. @overload
  133. def concatenate(
  134. *a: ArrayLike, axis: SupportsIndex = ..., out: _ArrayType = ...
  135. ) -> _ArrayType: ...
  136. @staticmethod
  137. def makemat(
  138. data: ArrayLike, dtype: DTypeLike = ..., copy: bool = ...
  139. ) -> _Matrix: ...
  140. # TODO: Sort out this `__getitem__` method
  141. def __getitem__(self, key: Any) -> Any: ...
  142. class RClass(AxisConcatenator):
  143. axis: Literal[0]
  144. matrix: Literal[False]
  145. ndmin: Literal[1]
  146. trans1d: Literal[-1]
  147. def __init__(self) -> None: ...
  148. r_: RClass
  149. class CClass(AxisConcatenator):
  150. axis: Literal[-1]
  151. matrix: Literal[False]
  152. ndmin: Literal[2]
  153. trans1d: Literal[0]
  154. def __init__(self) -> None: ...
  155. c_: CClass
  156. class IndexExpression(Generic[_BoolType]):
  157. maketuple: _BoolType
  158. def __init__(self, maketuple: _BoolType) -> None: ...
  159. @overload
  160. def __getitem__(self, item: _TupType) -> _TupType: ... # type: ignore[misc]
  161. @overload
  162. def __getitem__(self: IndexExpression[Literal[True]], item: _T) -> Tuple[_T]: ...
  163. @overload
  164. def __getitem__(self: IndexExpression[Literal[False]], item: _T) -> _T: ...
  165. index_exp: IndexExpression[Literal[True]]
  166. s_: IndexExpression[Literal[False]]
  167. def fill_diagonal(a: ndarray[Any, Any], val: Any, wrap: bool = ...) -> None: ...
  168. def diag_indices(n: int, ndim: int = ...) -> Tuple[NDArray[int_], ...]: ...
  169. def diag_indices_from(arr: ArrayLike) -> Tuple[NDArray[int_], ...]: ...
  170. # NOTE: see `numpy/__init__.pyi` for `ndenumerate` and `ndindex`