Skip to content
This repository was archived by the owner on Jun 10, 2020. It is now read-only.

MAINT: Introduced a couple of fixes for the np.core.fromnumeric functions #74

Merged
merged 9 commits into from
May 1, 2020
52 changes: 41 additions & 11 deletions numpy-stubs/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -813,14 +813,27 @@ _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]
# Integers and booleans can generally be used interchangeably
_ScalarInt = TypeVar("_ScalarInt", bound=Union[integer, bool_])
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think _ScalarInt is the best name here; seems like it should be _ScalarIntOrBool. I mostly don't want this to accidentally get cargo-culted to other places where NumPy int/bool diverge; e.g. arithmetic.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Based on #74 (comment) I'd argue that it is a suitable name. It can still be changed though if you don't agree,

Besides, for arithmetic it would (most likely) be more convenient to use a typevar of np.number anyway.

_ScalarGeneric = TypeVar("_ScalarGeneric", bound=generic)
_ScalarGenericPlus = TypeVar(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly _ScalarGenericPlus is not very descriptive of the type here. Also this definition would allow passing through subclasses of datetime and timedelta; are we confident that will work correctly?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe something like _ScalarGenericDT instead?

About your second concern, some quick testing suggests that subclassing does indeed work:

>>> import datetime as dt
>>> import numpy as np

>>> class A(dt.timedelta): ...
>>> np.take(A(1), 0).__class__
__main__.A

>>> class B(dt.datetime): ...
>>> np.take(B(2000, 1, 1), 0).__class__
__main__.B

"_ScalarGenericPlus", bound=Union[dt.datetime, dt.timedelta, generic]
)

# An array-like object consisting of integers
_Int = Union[int, integer]
_Int = Union[int, integer, bool, bool_]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As above about naming.

_Bool = Union[bool, bool_]
_ArrayLikeIntNested = Any # TODO: wait for support for recursive types
_ArrayLikeInt = Union[_Int, ndarray, Sequence[_Int], Sequence[_ArrayLikeIntNested]]
_ArrayLikeBoolNested = Any # TODO: wait for support for recursive types

# Integers and booleans can generally be used interchangeably
_ArrayLikeInt = Union[
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As above about naming.

_Int,
ndarray,
Sequence[_Int],
Sequence[_ArrayLikeIntNested],
Sequence[_ArrayLikeBoolNested],
]

# The signature of take() follows a common theme with its overloads:
# 1. A generic comes in; the same generic comes out
Expand All @@ -829,12 +842,12 @@ _ArrayLikeInt = Union[_Int, ndarray, Sequence[_Int], Sequence[_ArrayLikeIntNeste
# 4. An array-like object comes in; an ndarray or generic comes out
@overload
def take(
a: _ScalarGeneric,
a: _ScalarGenericPlus,
indices: int,
axis: Optional[int] = ...,
out: Optional[ndarray] = ...,
mode: _Mode = ...,
) -> _ScalarGeneric: ...
) -> _ScalarGenericPlus: ...
@overload
def take(
a: _Scalar,
Expand Down Expand Up @@ -862,21 +875,21 @@ def take(
def reshape(a: _ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ...
@overload
def choose(
a: _ScalarGeneric,
a: _ScalarInt,
choices: Union[Sequence[_ArrayLike], ndarray],
out: Optional[ndarray] = ...,
mode: _Mode = ...,
) -> _ScalarGeneric: ...
) -> _ScalarInt: ...
@overload
def choose(
a: _Scalar,
a: _Int,
choices: Union[Sequence[_ArrayLike], ndarray],
out: Optional[ndarray] = ...,
mode: _Mode = ...,
) -> _ScalarNumpy: ...
) -> Union[integer, bool_]: ...
@overload
def choose(
a: _ArrayLike,
a: _ArrayLikeInt,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So while passing bools technically works, given that the the docstring says

"This array must contain integers in [0, n-1], where n is the number of choices"

it seems to me that this is more an accident than the intent. I'm not sure that we should allow this in the types as it seems like not a best practice. @rgommers WDYT about this case? (We could also solicit opinions on the mailing list.)

Copy link
Member Author

@BvB93 BvB93 Apr 27, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Accident might not be the best term here.
The general compatibility between int and bool stems from the fact the bool is a subclass of the former.

>>> import numpy as np

>>> issubclass(int, bool)
True

>>> True == 1 and False == 0
True

>>> np.bool_(True) == 1 and np.bool_(False) == 0
True

Union[int, bool] is consequently equivalent to just plain int, but as np.bool_ is not a subclass of np.integer this does not hold for np.bool_. Even though the latter two classes behave (more or less; arithmetic is a bit of an exception here) identically to int and bool.

This also means that we could remove np.bool_ from the _ArrayLikeInt union, but properly removing bool would require adding a new overload along the lines of
choose(a: bool, ...) -> NoReturn.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

but as np.bool_ is not a subclass of np.integer this does not hold for np.bool_

And this is why I don't want to include it. Yes, including int gets you bool, but that's intentionally not true for numpy.bool_, so I'm not sure we should make an extra effort to bring it back here. To be clear, when I said "accident" I meant "that np.bool_ works".

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair enough, I've implemented the name changes in #74.

choices: Union[Sequence[_ArrayLike], ndarray],
out: Optional[ndarray] = ...,
mode: _Mode = ...,
Expand All @@ -898,6 +911,23 @@ def partition(
kind: _PartitionKind = ...,
order: Union[None, str, Sequence[str]] = ...,
) -> ndarray: ...
@overload
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The behaviour of np.argpartition() here is a bit odd.
If an np.ndarray is passed an np.ndarray is returned and if a np.generic is passed a np.generic is returned, so far so good.

However, if a builtin scalar is passed it also returns a np.ndarray, rather than a np.generic.

>>> import numpy as np

>>> type(np.argpartition(0, 0))
numpy.ndarray

>>> type(np.argpartition(np.int64(0), 0))
numpy.int64

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ouch. Sounds like it could potentially be worthy of opening an issue against NumPy.

def argpartition(
a: generic,
kth: _ArrayLikeInt,
axis: Optional[int] = ...,
kind: _PartitionKind = ...,
order: Union[None, str, Sequence[str]] = ...,
) -> integer: ...
@overload
def argpartition(
a: _ScalarBuiltin,
kth: _ArrayLikeInt,
axis: Optional[int] = ...,
kind: _PartitionKind = ...,
order: Union[None, str, Sequence[str]] = ...,
) -> ndarray: ...
@overload
def argpartition(
a: _ArrayLike,
kth: _ArrayLikeInt,
Expand Down
26 changes: 14 additions & 12 deletions tests/fail/fromnumeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,17 @@

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.take(a, None) # E: No overload variant of "take" matches argument type
np.take(a, axis=1.0) # E: No overload variant of "take" matches argument type
np.take(A, out=1) # E: No overload variant of "take" matches argument type
np.take(A, mode="bob") # E: No overload variant of "take" matches argument type

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.choose(a, None) # E: No overload variant of "choose" matches argument type
np.choose(a, out=1.0) # E: No overload variant of "choose" matches argument type
np.choose(A, mode="bob") # E: No overload variant of "choose" matches argument type

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
Expand All @@ -40,12 +40,14 @@
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( # E: No overload variant of "argpartition" matches argument type
a, None
)
np.argpartition(
A, 0, kind="bob" # E: Argument "kind" to "argpartition" has incompatible type
np.argpartition( # E: No overload variant of "argpartition" matches argument type
a, 0, axis="bob"
)
np.argpartition( # E: No overload variant of "argpartition" matches argument type
A, 0, kind="bob"
)
np.argpartition(
A, 0, order=range(5) # E: Argument "order" to "argpartition" has incompatible type
Expand Down
13 changes: 5 additions & 8 deletions tests/pass/fromnumeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,8 @@
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.choose(a, [True, True])
np.choose(A, [1.0, 1.0])

np.repeat(a, 1)
np.repeat(b, 1)
Expand All @@ -46,9 +43,9 @@
np.transpose(A)
np.transpose(B)

np.partition(a, 0)
np.partition(b, 0)
np.partition(c, 0)
np.partition(a, 0, axis=None)
np.partition(b, 0, axis=None)
np.partition(c, 0, axis=None)
np.partition(A, 0)
np.partition(B, 0)

Expand Down
21 changes: 7 additions & 14 deletions tests/reveal/fromnumeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,15 +39,8 @@
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.choose(a, [True, True])) # E: numpy.bool_
reveal_type(np.choose(A, [True, True])) # E: numpy.ndarray

reveal_type(np.repeat(a, 1)) # E: numpy.ndarray
reveal_type(np.repeat(b, 1)) # E: numpy.ndarray
Expand All @@ -66,14 +59,14 @@
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, axis=None)) # E: numpy.ndarray
reveal_type(np.partition(b, 0, axis=None)) # E: numpy.ndarray
reveal_type(np.partition(c, 0, axis=None)) # 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(a, 0)) # E: numpy.integer
reveal_type(np.argpartition(b, 0)) # E: numpy.integer
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
Expand Down