diff options
| author | Bas van Beek <b.f.van.beek@vu.nl> | 2021-05-21 14:11:49 +0200 |
|---|---|---|
| committer | Bas van Beek <43369155+BvB93@users.noreply.github.com> | 2021-06-06 20:54:37 +0200 |
| commit | 3fe619717c7031ad089adc289fbd7a5f61ea957a (patch) | |
| tree | 2432c8ff4ac7155d54b300a385f865de429c21a5 /numpy | |
| parent | 1dc53a3be6c4191bc65ad7a27ee571b07e52ee97 (diff) | |
| download | numpy-3fe619717c7031ad089adc289fbd7a5f61ea957a.tar.gz | |
MAINT: Move `ravel_multi_index` and `unravel_index` to `np.core.multiarray`
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/multiarray.pyi | 58 | ||||
| -rw-r--r-- | numpy/lib/index_tricks.pyi | 33 |
2 files changed, 60 insertions, 31 deletions
diff --git a/numpy/core/multiarray.pyi b/numpy/core/multiarray.pyi index 0080dc303..96f9edaa3 100644 --- a/numpy/core/multiarray.pyi +++ b/numpy/core/multiarray.pyi @@ -1,7 +1,31 @@ -from typing import Any, Optional, overload, TypeVar, List, Type, Union +import sys +from typing import ( + Any, + Optional, + overload, + TypeVar, + List, + Type, + Union, + Sequence, + Tuple, +) -from numpy import dtype, generic, _OrderKACF -from numpy.typing import ArrayLike, NDArray, DTypeLike, _ShapeLike, _SupportsDType +from numpy import dtype, generic, intp, _OrderKACF, _OrderCF, _ModeKind +from numpy.typing import ( + ArrayLike, + NDArray, + DTypeLike, + _SupportsDType, + _ShapeLike, + _IntLike_co, + _ArrayLikeInt_co, +) + +if sys.version_info >= (3, 8): + from typing import SupportsIndex +else: + from typing_extensions import SupportsIndex _SCT = TypeVar("_SCT", bound=generic) @@ -86,3 +110,31 @@ def empty( *, like: ArrayLike = ..., ) -> NDArray[Any]: ... + +@overload +def unravel_index( # type: ignore[misc] + indices: _IntLike_co, + shape: _ShapeLike, + order: _OrderCF = ..., +) -> Tuple[intp, ...]: ... +@overload +def unravel_index( + indices: _ArrayLikeInt_co, + shape: _ShapeLike, + order: _OrderCF = ..., +) -> Tuple[NDArray[intp], ...]: ... + +@overload +def ravel_multi_index( # type: ignore[misc] + multi_index: Sequence[_IntLike_co], + dims: Sequence[SupportsIndex], + mode: Union[_ModeKind, Tuple[_ModeKind, ...]] = ..., + order: _OrderCF = ..., +) -> intp: ... +@overload +def ravel_multi_index( + multi_index: Sequence[_ArrayLikeInt_co], + dims: Sequence[SupportsIndex], + mode: Union[_ModeKind, Tuple[_ModeKind, ...]] = ..., + order: _OrderCF = ..., +) -> NDArray[intp]: ... diff --git a/numpy/lib/index_tricks.pyi b/numpy/lib/index_tricks.pyi index a3bfef6b6..0f9ae94a9 100644 --- a/numpy/lib/index_tricks.pyi +++ b/numpy/lib/index_tricks.pyi @@ -44,6 +44,11 @@ from numpy.typing import ( _ShapeLike, ) +from numpy.core.multiarray import ( + unravel_index as unravel_index, + ravel_multi_index as ravel_multi_index, +) + if sys.version_info >= (3, 8): from typing import Literal, SupportsIndex else: @@ -58,34 +63,6 @@ _ArrayType = TypeVar("_ArrayType", bound=ndarray[Any, Any]) __all__: List[str] @overload -def unravel_index( # type: ignore[misc] - indices: Union[int, integer[Any]], - shape: _ShapeLike, - order: _OrderCF = ... -) -> Tuple[intp, ...]: ... -@overload -def unravel_index( - indices: _ArrayLikeInt, - shape: _ShapeLike, - order: _OrderCF = ... -) -> Tuple[NDArray[intp], ...]: ... - -@overload -def ravel_multi_index( # type: ignore[misc] - multi_index: Sequence[Union[int, integer[Any]]], - dims: _ShapeLike, - mode: Union[_ModeKind, Tuple[_ModeKind, ...]] = ..., - order: _OrderCF = ... -) -> intp: ... -@overload -def ravel_multi_index( - multi_index: Sequence[_ArrayLikeInt], - dims: _ShapeLike, - mode: Union[_ModeKind, Tuple[_ModeKind, ...]] = ..., - order: _OrderCF = ... -) -> NDArray[intp]: ... - -@overload def ix_(*args: _NestedSequence[_SupportsDType[_DType]]) -> Tuple[ndarray[Any, _DType], ...]: ... @overload def ix_(*args: _NestedSequence[str]) -> Tuple[NDArray[str_], ...]: ... |