From dd7a60119393c9c4c5f61a57b8f4037f598b1a1f Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Tue, 21 Apr 2020 13:55:42 +0200 Subject: [PATCH 1/8] Added ".DS_Store" to .gitgnore --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 01c1e46..ed2d84e 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,6 @@ venv .idea *~ **~ + +# MacOS +.DS_Store From 36581cbbfdb6dbae3235706bb4833d0f6f046183 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Tue, 21 Apr 2020 13:57:35 +0200 Subject: [PATCH 2/8] Added 11 functions from np.core.fromnumeric --- numpy-stubs/__init__.pyi | 126 ++++++++++++++++++++++++++++++++++++ tests/fail/fromnumeric.py | 60 +++++++++++++++++ tests/pass/fromnumeric.py | 63 ++++++++++++++++++ tests/reveal/fromnumeric.py | 65 +++++++++++++++++++ 4 files changed, 314 insertions(+) create mode 100644 tests/fail/fromnumeric.py create mode 100644 tests/pass/fromnumeric.py create mode 100644 tests/reveal/fromnumeric.py diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index 49a085a..ea4df8a 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,124 @@ 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"] + +_Generic = TypeVar("_Generic", bound=generic) + +# Can't use the SupportsInt protocol here as the ndarray.__int__() method exists +_ScalarBuiltin = Union[str, bytes, dt.date, dt.timedelta, bool, int, float, complex] +_Scalar = Union[_ScalarBuiltin, generic] + +# An array-like object consisting of integers +_ArrayLikeInt = Union[ + int, integer, Sequence[Union[int, integer]], ndarray +] # TODO: ndarray[int] + +# An array-like object consisting of strings +_ArrayLikeStr = Union[ + str, str_, Sequence[Union[str, str_]], ndarray +] # TODO: ndarray[str] + +# 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: _Generic, + indices: int, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> _Generic: ... +@overload +def take( + a: _Scalar, + indices: int, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> generic: ... +@overload +def take( + a: _ArrayLike, + indices: int, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> generic: ... +@overload +def take( + a: _ArrayLike, + indices: _ArrayLikeInt, + axis: Optional[int] = ..., + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> Union[generic, ndarray]: ... +def reshape(a: _ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ... +@overload +def choose( + a: _Generic, + choices: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike] + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> _Generic: ... +@overload +def choose( + a: _Scalar, + choices: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike] + out: Optional[ndarray] = ..., + mode: _Mode = ..., +) -> generic: ... +@overload +def choose( + a: _ArrayLike, + choices: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike] + 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], # TODO: ndarray[_ArrayLike] + axis1: int, + axis2: int, +) -> ndarray: ... +def transpose( + a: _ArrayLike, + axes: Union[None, Sequence[int], ndarray] = ..., # TODO: ndarray[int] +) -> ndarray: ... +def partition( + a: _ArrayLike, + kth: _ArrayLikeInt, + axis: Optional[int] = ..., + kind: _PartitionKind = ..., + order: Optional[_ArrayLikeStr] = ..., +) -> ndarray: ... +def argpartition( + a: _ArrayLike, + kth: _ArrayLikeInt, + axis: Optional[int] = ..., + kind: _PartitionKind = ..., + order: Optional[_ArrayLikeStr] = ..., +) -> ndarray: ... +def sort( + a: Union[Sequence[_ArrayLike], ndarray], + axis: Optional[int] = ..., + kind: Optional[_SortKind] = ..., + order: Optional[_ArrayLikeStr] = ..., +) -> ndarray: ... +def argsort( + a: Union[Sequence[_ArrayLike], ndarray], + axis: Optional[int] = ..., + kind: Optional[_SortKind] = ..., + order: Optional[_ArrayLikeStr] = ..., +) -> ndarray: ... diff --git a/tests/fail/fromnumeric.py b/tests/fail/fromnumeric.py new file mode 100644 index 0000000..56da171 --- /dev/null +++ b/tests/fail/fromnumeric.py @@ -0,0 +1,60 @@ +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..93b6f77 --- /dev/null +++ b/tests/pass/fromnumeric.py @@ -0,0 +1,63 @@ +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..3caa4e1 --- /dev/null +++ b/tests/reveal/fromnumeric.py @@ -0,0 +1,65 @@ +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: numpy.generic +reveal_type(np.take(A, 0)) # E: numpy.generic +reveal_type(np.take(B, 0)) # E: numpy.generic +reveal_type(np.take(A, [0])) # E: Union[numpy.generic, numpy.ndarray] +reveal_type(np.take(B, [0])) # E: Union[numpy.generic, numpy.ndarray] + +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(c, [1.0])) # E: numpy.generic +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 From c2eeada1bc1bdcb785993acadb7aacf4b738a795 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Tue, 21 Apr 2020 14:20:30 +0200 Subject: [PATCH 3/8] Downgraded black from 19.10b0 to 18.9b0 --- numpy-stubs/__init__.pyi | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index ea4df8a..f4b0759 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -889,8 +889,7 @@ def swapaxes( axis2: int, ) -> ndarray: ... def transpose( - a: _ArrayLike, - axes: Union[None, Sequence[int], ndarray] = ..., # TODO: ndarray[int] + a: _ArrayLike, axes: Union[None, Sequence[int], ndarray] = ... # TODO: ndarray[int] ) -> ndarray: ... def partition( a: _ArrayLike, From 37cbdedd2a60bbb92c1eb3b52b65a4c6786058ec Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Tue, 21 Apr 2020 20:27:44 +0200 Subject: [PATCH 4/8] Added a module-level docstring with the name of the corresponding numpy module --- tests/fail/fromnumeric.py | 2 ++ tests/pass/fromnumeric.py | 2 ++ tests/reveal/fromnumeric.py | 2 ++ 3 files changed, 6 insertions(+) diff --git a/tests/fail/fromnumeric.py b/tests/fail/fromnumeric.py index 56da171..9d01f21 100644 --- a/tests/fail/fromnumeric.py +++ b/tests/fail/fromnumeric.py @@ -1,3 +1,5 @@ +"""Tests for :mod:`numpy.core.fromnumeric`.""" + import numpy as np A = np.array(True, ndmin=2, dtype=bool) diff --git a/tests/pass/fromnumeric.py b/tests/pass/fromnumeric.py index 93b6f77..1e7ba57 100644 --- a/tests/pass/fromnumeric.py +++ b/tests/pass/fromnumeric.py @@ -1,3 +1,5 @@ +"""Tests for :mod:`numpy.core.fromnumeric`.""" + import numpy as np A = np.array(True, ndmin=2, dtype=bool) diff --git a/tests/reveal/fromnumeric.py b/tests/reveal/fromnumeric.py index 3caa4e1..7394678 100644 --- a/tests/reveal/fromnumeric.py +++ b/tests/reveal/fromnumeric.py @@ -1,3 +1,5 @@ +"""Tests for :mod:`numpy.core.fromnumeric`.""" + import numpy as np A = np.array(True, ndmin=2, dtype=bool) From 36efcf359a24bb2d9a448bf8b99738221405bca5 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Tue, 21 Apr 2020 20:32:15 +0200 Subject: [PATCH 5/8] Implemented feedback from https://github.com/numpy/numpy-stubs/pull/67#pullrequestreview-397431275 * Renamed the ``_Scalar`` TypeVar to ``_ScalarGeneric`` (https://github.com/numpy/numpy-stubs/pull/67#discussion_r412273202) * Expanded the scope of ``_ArrayLikeInt`` up to sextuple levels of Sequence nesting (https://github.com/numpy/numpy-stubs/pull/67#discussion_r412279088 and https://github.com/numpy/numpy-stubs/pull/67#discussion_r412289943) --- numpy-stubs/__init__.pyi | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index f4b0759..bddecbf 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -804,16 +804,24 @@ _Order = Literal["C", "F", "A"] _PartitionKind = Literal["introselect"] _SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"] -_Generic = TypeVar("_Generic", bound=generic) - -# Can't use the SupportsInt protocol here as the ndarray.__int__() method exists +# Various annotations for scalars +_ScalarGeneric = TypeVar("_ScalarGeneric", bound=generic) _ScalarBuiltin = Union[str, bytes, dt.date, dt.timedelta, bool, int, float, complex] _Scalar = Union[_ScalarBuiltin, generic] # An array-like object consisting of integers +# TODO: If possible, figure out a better way to deal with nested sequences +_Int = Union[int, integer] _ArrayLikeInt = Union[ - int, integer, Sequence[Union[int, integer]], ndarray -] # TODO: ndarray[int] + _Int, + ndarray, # TODO: ndarray[int] + Sequence[_Int], + Sequence[Sequence[_Int]], + Sequence[Sequence[Sequence[_Int]]], + Sequence[Sequence[Sequence[Sequence[_Int]]]], + Sequence[Sequence[Sequence[Sequence[Sequence[_Int]]]]], + Sequence[Sequence[Sequence[Sequence[Sequence[Sequence[_Int]]]]]], +] # An array-like object consisting of strings _ArrayLikeStr = Union[ @@ -827,12 +835,12 @@ _ArrayLikeStr = Union[ # 4. An array-like object comes in; an ndarray or generic comes out @overload def take( - a: _Generic, + a: _ScalarGeneric, indices: int, axis: Optional[int] = ..., out: Optional[ndarray] = ..., mode: _Mode = ..., -) -> _Generic: ... +) -> _ScalarGeneric: ... @overload def take( a: _Scalar, @@ -860,11 +868,11 @@ def take( def reshape(a: _ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ... @overload def choose( - a: _Generic, + a: _ScalarGeneric, choices: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike] out: Optional[ndarray] = ..., mode: _Mode = ..., -) -> _Generic: ... +) -> _ScalarGeneric: ... @overload def choose( a: _Scalar, From 837eec640b4cca8616f316e0542f0cfd025ab4a5 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Tue, 21 Apr 2020 21:24:47 +0200 Subject: [PATCH 6/8] Replaced the nested unions with _ArrayLikeIntNested --- numpy-stubs/__init__.pyi | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index bddecbf..3d8b99a 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -810,17 +810,13 @@ _ScalarBuiltin = Union[str, bytes, dt.date, dt.timedelta, bool, int, float, comp _Scalar = Union[_ScalarBuiltin, generic] # An array-like object consisting of integers -# TODO: If possible, figure out a better way to deal with nested sequences _Int = Union[int, integer] +_ArrayLikeIntNested = Any # TODO: wait for support for recursive types _ArrayLikeInt = Union[ _Int, ndarray, # TODO: ndarray[int] Sequence[_Int], - Sequence[Sequence[_Int]], - Sequence[Sequence[Sequence[_Int]]], - Sequence[Sequence[Sequence[Sequence[_Int]]]], - Sequence[Sequence[Sequence[Sequence[Sequence[_Int]]]]], - Sequence[Sequence[Sequence[Sequence[Sequence[Sequence[_Int]]]]]], + Sequence[_ArrayLikeIntNested] ] # An array-like object consisting of strings From 11dd7e4ba92470fd908747829d82e33654f95921 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Tue, 21 Apr 2020 21:50:19 +0200 Subject: [PATCH 7/8] Ran black again --- numpy-stubs/__init__.pyi | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index 3d8b99a..643758d 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -813,10 +813,7 @@ _Scalar = Union[_ScalarBuiltin, generic] _Int = Union[int, integer] _ArrayLikeIntNested = Any # TODO: wait for support for recursive types _ArrayLikeInt = Union[ - _Int, - ndarray, # TODO: ndarray[int] - Sequence[_Int], - Sequence[_ArrayLikeIntNested] + _Int, ndarray, Sequence[_Int], Sequence[_ArrayLikeIntNested] # TODO: ndarray[int] ] # An array-like object consisting of strings From 35e1cdade32d8d704599a63f0a1545abc6596f63 Mon Sep 17 00:00:00 2001 From: Bas van Beek Date: Wed, 22 Apr 2020 18:17:50 +0200 Subject: [PATCH 8/8] Implement the feedback from https://github.com/numpy/numpy-stubs/pull/67#pullrequestreview-397782543 * The ``order`` parameter cannot be an ndarray, it has to be proper Sequence; change its annotation to ``Union[None, str, Sequence[str]]`` (see https://github.com/numpy/numpy-stubs/pull/67#discussion_r412607079, https://github.com/numpy/numpy-stubs/pull/67#discussion_r412607730 and https://github.com/numpy/numpy-stubs/pull/67#discussion_r412607869) * ``datetime.datetime`` and ``datetime.timedelta`` are also valid return types besides ``np.generic`` and ``np.ndarray`` (see https://github.com/numpy/numpy-stubs/pull/67#discussion_r412604113) * Removed all ``"TODO"`` comments expect for ``_ArrayLikeIntNested`` (https://github.com/numpy/numpy-stubs/pull/67#pullrequestreview-397782543) --- numpy-stubs/__init__.pyi | 49 ++++++++++++++++++------------------- tests/reveal/fromnumeric.py | 30 ++++++++++++++++++----- 2 files changed, 48 insertions(+), 31 deletions(-) diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index 643758d..c51f3a6 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -805,21 +805,22 @@ _PartitionKind = Literal["introselect"] _SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"] # Various annotations for scalars -_ScalarGeneric = TypeVar("_ScalarGeneric", bound=generic) + +# 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, generic] +_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] # TODO: ndarray[int] -] - -# An array-like object consisting of strings -_ArrayLikeStr = Union[ - str, str_, Sequence[Union[str, str_]], ndarray -] # TODO: ndarray[str] +_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 @@ -841,7 +842,7 @@ def take( axis: Optional[int] = ..., out: Optional[ndarray] = ..., mode: _Mode = ..., -) -> generic: ... +) -> _ScalarNumpy: ... @overload def take( a: _ArrayLike, @@ -849,7 +850,7 @@ def take( axis: Optional[int] = ..., out: Optional[ndarray] = ..., mode: _Mode = ..., -) -> generic: ... +) -> _ScalarNumpy: ... @overload def take( a: _ArrayLike, @@ -857,26 +858,26 @@ def take( axis: Optional[int] = ..., out: Optional[ndarray] = ..., mode: _Mode = ..., -) -> Union[generic, ndarray]: ... +) -> Union[_ScalarNumpy, ndarray]: ... def reshape(a: _ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ... @overload def choose( a: _ScalarGeneric, - choices: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike] + choices: Union[Sequence[_ArrayLike], ndarray], out: Optional[ndarray] = ..., mode: _Mode = ..., ) -> _ScalarGeneric: ... @overload def choose( a: _Scalar, - choices: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike] + choices: Union[Sequence[_ArrayLike], ndarray], out: Optional[ndarray] = ..., mode: _Mode = ..., -) -> generic: ... +) -> _ScalarNumpy: ... @overload def choose( a: _ArrayLike, - choices: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike] + choices: Union[Sequence[_ArrayLike], ndarray], out: Optional[ndarray] = ..., mode: _Mode = ..., ) -> ndarray: ... @@ -885,36 +886,34 @@ def repeat( ) -> ndarray: ... def put(a: ndarray, ind: _ArrayLikeInt, v: _ArrayLike, mode: _Mode = ...) -> None: ... def swapaxes( - a: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike] - axis1: int, - axis2: int, + a: Union[Sequence[_ArrayLike], ndarray], axis1: int, axis2: int ) -> ndarray: ... def transpose( - a: _ArrayLike, axes: Union[None, Sequence[int], ndarray] = ... # TODO: ndarray[int] + a: _ArrayLike, axes: Union[None, Sequence[int], ndarray] = ... ) -> ndarray: ... def partition( a: _ArrayLike, kth: _ArrayLikeInt, axis: Optional[int] = ..., kind: _PartitionKind = ..., - order: Optional[_ArrayLikeStr] = ..., + order: Union[None, str, Sequence[str]] = ..., ) -> ndarray: ... def argpartition( a: _ArrayLike, kth: _ArrayLikeInt, axis: Optional[int] = ..., kind: _PartitionKind = ..., - order: Optional[_ArrayLikeStr] = ..., + order: Union[None, str, Sequence[str]] = ..., ) -> ndarray: ... def sort( a: Union[Sequence[_ArrayLike], ndarray], axis: Optional[int] = ..., kind: Optional[_SortKind] = ..., - order: Optional[_ArrayLikeStr] = ..., + order: Union[None, str, Sequence[str]] = ..., ) -> ndarray: ... def argsort( a: Union[Sequence[_ArrayLike], ndarray], axis: Optional[int] = ..., kind: Optional[_SortKind] = ..., - order: Optional[_ArrayLikeStr] = ..., + order: Union[None, str, Sequence[str]] = ..., ) -> ndarray: ... diff --git a/tests/reveal/fromnumeric.py b/tests/reveal/fromnumeric.py index 7394678..82450df 100644 --- a/tests/reveal/fromnumeric.py +++ b/tests/reveal/fromnumeric.py @@ -13,11 +13,25 @@ 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: numpy.generic -reveal_type(np.take(A, 0)) # E: numpy.generic -reveal_type(np.take(B, 0)) # E: numpy.generic -reveal_type(np.take(A, [0])) # E: Union[numpy.generic, numpy.ndarray] -reveal_type(np.take(B, [0])) # E: Union[numpy.generic, numpy.ndarray] +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 @@ -27,7 +41,11 @@ reveal_type(np.choose(a, [True])) # E: numpy.bool_ reveal_type(np.choose(b, [1.0])) # E: numpy.float32 -reveal_type(np.choose(c, [1.0])) # E: numpy.generic +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