summaryrefslogtreecommitdiff
path: root/numpy/core
diff options
context:
space:
mode:
Diffstat (limited to 'numpy/core')
-rw-r--r--numpy/core/multiarray.pyi58
1 files changed, 55 insertions, 3 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]: ...