diff --git a/.gitignore b/.gitignore index 01c1e46..ed2d84e 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,6 @@ venv .idea *~ **~ + +# MacOS +.DS_Store diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index 49a085a..c51f3a6 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -35,6 +35,11 @@ if sys.version_info[0] < 3: else: from typing import SupportsBytes +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + # TODO: remove when the full numpy namespace is defined def __getattr__(name: str) -> Any: ... @@ -792,3 +797,123 @@ class AxisError(ValueError, IndexError): def __init__( self, axis: int, ndim: Optional[int] = ..., msg_prefix: Optional[str] = ... ) -> None: ... + +# Functions from np.core.fromnumeric +_Mode = Literal["raise", "wrap", "clip"] +_Order = Literal["C", "F", "A"] +_PartitionKind = Literal["introselect"] +_SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"] + +# Various annotations for scalars + +# While dt.datetime and dt.timedelta are not technically part of NumPy, +# they are one of the rare few builtin scalars which serve as valid return types. +# See https://github.com/numpy/numpy-stubs/pull/67#discussion_r412604113. +_ScalarNumpy = Union[generic, dt.datetime, dt.timedelta] +_ScalarBuiltin = Union[str, bytes, dt.date, dt.timedelta, bool, int, float, complex] +_Scalar = Union[_ScalarBuiltin, _ScalarNumpy] + +_ScalarGeneric = TypeVar( + "_ScalarGeneric", bound=Union[dt.datetime, dt.timedelta, generic] +) + +# An array-like object consisting of integers +_Int = Union[int, integer] +_ArrayLikeIntNested = Any # TODO: wait for support for recursive types +_ArrayLikeInt = Union[_Int, ndarray, Sequence[_Int], Sequence[_ArrayLikeIntNested]] + +# The signature of take() follows a common theme with its overloads: +# 1. A generic comes in; the same generic comes out +# 2. A scalar comes in; a generic comes out +# 3. An array-like object comes in; some keyword ensures that a generic comes out +# 4. An array-like object comes in; an ndarray or generic comes out +@overload +def take( + a: _ScalarGeneric, + indices: int, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> _ScalarGeneric: ... +@overload +def take( + a: _Scalar, + indices: int, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> _ScalarNumpy: ... +@overload +def take( + a: _ArrayLike, + indices: int, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> _ScalarNumpy: ... +@overload +def take( + a: _ArrayLike, + indices: _ArrayLikeInt, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> Union[_ScalarNumpy, ndarray]: ... +def reshape(a: _ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ... +@overload +def choose( + a: _ScalarGeneric, + choices: Union[Sequence[_ArrayLike], ndarray], + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> _ScalarGeneric: ... +@overload +def choose( + a: _Scalar, + choices: Union[Sequence[_ArrayLike], ndarray], + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> _ScalarNumpy: ... +@overload +def choose( + a: _ArrayLike, + choices: Union[Sequence[_ArrayLike], ndarray], + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> ndarray: ... +def repeat( + a: _ArrayLike, repeats: _ArrayLikeInt, axis: Optional[int] = ... +) -> ndarray: ... +def put(a: ndarray, ind: _ArrayLikeInt, v: _ArrayLike, mode: _Mode = ...) -> None: ... +def swapaxes( + a: Union[Sequence[_ArrayLike], ndarray], axis1: int, axis2: int +) -> ndarray: ... +def transpose( + a: _ArrayLike, axes: Union[None, Sequence[int], ndarray] = ... +) -> ndarray: ... +def partition( + a: _ArrayLike, + kth: _ArrayLikeInt, + axis: Optional[int] = ..., + kind: _PartitionKind = ..., + order: Union[None, str, Sequence[str]] = ..., +) -> ndarray: ... +def argpartition( + a: _ArrayLike, + kth: _ArrayLikeInt, + axis: Optional[int] = ..., + kind: _PartitionKind = ..., + order: Union[None, str, Sequence[str]] = ..., +) -> ndarray: ... +def sort( + a: Union[Sequence[_ArrayLike], ndarray], + axis: Optional[int] = ..., + kind: Optional[_SortKind] = ..., + order: Union[None, str, Sequence[str]] = ..., +) -> ndarray: ... +def argsort( + a: Union[Sequence[_ArrayLike], ndarray], + axis: Optional[int] = ..., + kind: Optional[_SortKind] = ..., + order: Union[None, str, Sequence[str]] = ..., +) -> ndarray: ... diff --git a/tests/fail/fromnumeric.py b/tests/fail/fromnumeric.py new file mode 100644 index 0000000..9d01f21 --- /dev/null +++ b/tests/fail/fromnumeric.py @@ -0,0 +1,62 @@ +"""Tests for :mod:`numpy.core.fromnumeric`.""" + +import numpy as np + +A = np.array(True, ndmin=2, dtype=bool) +A.setflags(write=False) + +a = np.bool_(True) + +np.take(a, None) # E: No overload variant of "take" matches argument types +np.take(a, axis=1.0) # E: No overload variant of "take" matches argument types +np.take(A, out=1) # E: No overload variant of "take" matches argument types +np.take(A, mode="bob") # E: No overload variant of "take" matches argument types + +np.reshape(a, None) # E: Argument 2 to "reshape" has incompatible type +np.reshape(A, 1, order="bob") # E: Argument "order" to "reshape" has incompatible type + +np.choose(a, None) # E: No overload variant of "choose" matches argument types +np.choose(a, out=1.0) # E: No overload variant of "choose" matches argument types +np.choose(A, mode="bob") # E: No overload variant of "choose" matches argument types + +np.repeat(a, None) # E: Argument 2 to "repeat" has incompatible type +np.repeat(A, 1, axis=1.0) # E: Argument "axis" to "repeat" has incompatible type + +np.swapaxes(a, 0, 0) # E: Argument 1 to "swapaxes" has incompatible type +np.swapaxes(A, None, 1) # E: Argument 2 to "swapaxes" has incompatible type +np.swapaxes(A, 1, [0]) # E: Argument 3 to "swapaxes" has incompatible type + +np.transpose(a, axes=1) # E: Argument "axes" to "transpose" has incompatible type +np.transpose(A, axes=1.0) # E: Argument "axes" to "transpose" has incompatible type + +np.partition(a, None) # E: Argument 2 to "partition" has incompatible type +np.partition( + a, 0, axis="bob" # E: Argument "axis" to "partition" has incompatible type +) +np.partition( + A, 0, kind="bob" # E: Argument "kind" to "partition" has incompatible type +) +np.partition( + A, 0, order=range(5) # E: Argument "order" to "partition" has incompatible type +) + +np.argpartition(a, None) # E: Argument 2 to "argpartition" has incompatible type +np.argpartition( + a, 0, axis="bob" # E: Argument "axis" to "argpartition" has incompatible type +) +np.argpartition( + A, 0, kind="bob" # E: Argument "kind" to "argpartition" has incompatible type +) +np.argpartition( + A, 0, order=range(5) # E: Argument "order" to "argpartition" has incompatible type +) + +np.sort(a) # E: Argument 1 to "sort" has incompatible type +np.sort(A, axis="bob") # E: Argument "axis" to "sort" has incompatible type +np.sort(A, kind="bob") # E: Argument "kind" to "sort" has incompatible type +np.sort(A, order=range(5)) # E: Argument "order" to "sort" has incompatible type + +np.argsort(a) # E: Argument 1 to "argsort" has incompatible type +np.argsort(A, axis="bob") # E: Argument "axis" to "argsort" has incompatible type +np.argsort(A, kind="bob") # E: Argument "kind" to "argsort" has incompatible type +np.argsort(A, order=range(5)) # E: Argument "order" to "argsort" has incompatible type diff --git a/tests/pass/fromnumeric.py b/tests/pass/fromnumeric.py new file mode 100644 index 0000000..1e7ba57 --- /dev/null +++ b/tests/pass/fromnumeric.py @@ -0,0 +1,65 @@ +"""Tests for :mod:`numpy.core.fromnumeric`.""" + +import numpy as np + +A = np.array(True, ndmin=2, dtype=bool) +B = np.array(1.0, ndmin=2, dtype=np.float32) +A.setflags(write=False) +B.setflags(write=False) + +a = np.bool_(True) +b = np.float32(1.0) +c = 1.0 + +np.take(a, 0) +np.take(b, 0) +np.take(c, 0) +np.take(A, 0) +np.take(B, 0) +np.take(A, [0]) +np.take(B, [0]) + +np.reshape(a, 1) +np.reshape(b, 1) +np.reshape(c, 1) +np.reshape(A, 1) +np.reshape(B, 1) + +np.choose(a, [True]) +np.choose(b, [1.0]) +np.choose(c, [1.0]) +np.choose(A, [True]) +np.choose(B, [1.0]) + +np.repeat(a, 1) +np.repeat(b, 1) +np.repeat(c, 1) +np.repeat(A, 1) +np.repeat(B, 1) + +np.swapaxes(A, 0, 0) +np.swapaxes(B, 0, 0) + +np.transpose(a) +np.transpose(b) +np.transpose(c) +np.transpose(A) +np.transpose(B) + +np.partition(a, 0) +np.partition(b, 0) +np.partition(c, 0) +np.partition(A, 0) +np.partition(B, 0) + +np.argpartition(a, 0) +np.argpartition(b, 0) +np.argpartition(c, 0) +np.argpartition(A, 0) +np.argpartition(B, 0) + +np.sort(A, 0) +np.sort(B, 0) + +np.argsort(A, 0) +np.argsort(B, 0) diff --git a/tests/reveal/fromnumeric.py b/tests/reveal/fromnumeric.py new file mode 100644 index 0000000..82450df --- /dev/null +++ b/tests/reveal/fromnumeric.py @@ -0,0 +1,85 @@ +"""Tests for :mod:`numpy.core.fromnumeric`.""" + +import numpy as np + +A = np.array(True, ndmin=2, dtype=bool) +B = np.array(1.0, ndmin=2, dtype=np.float32) +A.setflags(write=False) +B.setflags(write=False) + +a = np.bool_(True) +b = np.float32(1.0) +c = 1.0 + +reveal_type(np.take(a, 0)) # E: numpy.bool_ +reveal_type(np.take(b, 0)) # E: numpy.float32 +reveal_type( + np.take(c, 0) # E: Union[numpy.generic, datetime.datetime, datetime.timedelta] +) +reveal_type( + np.take(A, 0) # E: Union[numpy.generic, datetime.datetime, datetime.timedelta] +) +reveal_type( + np.take(B, 0) # E: Union[numpy.generic, datetime.datetime, datetime.timedelta] +) +reveal_type( + np.take( # E: Union[numpy.generic, datetime.datetime, datetime.timedelta, numpy.ndarray] + A, [0] + ) +) +reveal_type( + np.take( # E: Union[numpy.generic, datetime.datetime, datetime.timedelta, numpy.ndarray] + B, [0] + ) +) + +reveal_type(np.reshape(a, 1)) # E: numpy.ndarray +reveal_type(np.reshape(b, 1)) # E: numpy.ndarray +reveal_type(np.reshape(c, 1)) # E: numpy.ndarray +reveal_type(np.reshape(A, 1)) # E: numpy.ndarray +reveal_type(np.reshape(B, 1)) # E: numpy.ndarray + +reveal_type(np.choose(a, [True])) # E: numpy.bool_ +reveal_type(np.choose(b, [1.0])) # E: numpy.float32 +reveal_type( + np.choose( # E: Union[numpy.generic, datetime.datetime, datetime.timedelta] + c, [1.0] + ) +) +reveal_type(np.choose(A, [True])) # E: numpy.ndarray +reveal_type(np.choose(B, [1.0])) # E: numpy.ndarray + +reveal_type(np.repeat(a, 1)) # E: numpy.ndarray +reveal_type(np.repeat(b, 1)) # E: numpy.ndarray +reveal_type(np.repeat(c, 1)) # E: numpy.ndarray +reveal_type(np.repeat(A, 1)) # E: numpy.ndarray +reveal_type(np.repeat(B, 1)) # E: numpy.ndarray + +# TODO: Add tests for np.put() + +reveal_type(np.swapaxes(A, 0, 0)) # E: numpy.ndarray +reveal_type(np.swapaxes(B, 0, 0)) # E: numpy.ndarray + +reveal_type(np.transpose(a)) # E: numpy.ndarray +reveal_type(np.transpose(b)) # E: numpy.ndarray +reveal_type(np.transpose(c)) # E: numpy.ndarray +reveal_type(np.transpose(A)) # E: numpy.ndarray +reveal_type(np.transpose(B)) # E: numpy.ndarray + +reveal_type(np.partition(a, 0)) # E: numpy.ndarray +reveal_type(np.partition(b, 0)) # E: numpy.ndarray +reveal_type(np.partition(c, 0)) # E: numpy.ndarray +reveal_type(np.partition(A, 0)) # E: numpy.ndarray +reveal_type(np.partition(B, 0)) # E: numpy.ndarray + +reveal_type(np.argpartition(a, 0)) # E: numpy.ndarray +reveal_type(np.argpartition(b, 0)) # E: numpy.ndarray +reveal_type(np.argpartition(c, 0)) # E: numpy.ndarray +reveal_type(np.argpartition(A, 0)) # E: numpy.ndarray +reveal_type(np.argpartition(B, 0)) # E: numpy.ndarray + +reveal_type(np.sort(A, 0)) # E: numpy.ndarray +reveal_type(np.sort(B, 0)) # E: numpy.ndarray + +reveal_type(np.argsort(A, 0)) # E: numpy.ndarray +reveal_type(np.argsort(B, 0)) # E: numpy.ndarray