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

ENH: Add type annotations for the np.core.fromnumeric module: part 1/4 #67

Merged
merged 8 commits into from
Apr 24, 2020
Merged

ENH: Add type annotations for the np.core.fromnumeric module: part 1/4 #67

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
dd7a601
Added ".DS_Store" to .gitgnore
Apr 21, 2020
36581cb
Added 11 functions from np.core.fromnumeric
Apr 21, 2020
c2eeada
Downgraded black from 19.10b0 to 18.9b0
Apr 21, 2020
37cbded
Added a module-level docstring with the name of the corresponding num…
Apr 21, 2020
36efcf3
Implemented feedback from https://github.com/numpy/numpy-stubs/pull/6…
Apr 21, 2020
837eec6
Replaced the nested unions with _ArrayLikeIntNested
Apr 21, 2020
11dd7e4
Ran black again
Apr 21, 2020
35e1cda
Implement the feedback from https://github.com/numpy/numpy-stubs/pull…
Apr 22, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,6 @@ venv
.idea
*~
**~

# MacOS
.DS_Store
133 changes: 133 additions & 0 deletions numpy-stubs/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...

Expand Down Expand Up @@ -792,3 +797,131 @@ 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
_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,
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[
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: _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 = ...,
) -> generic: ...
Copy link
Member

Choose a reason for hiding this comment

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

This isn't always true I don't think:

>>> type(np.take(datetime.timedelta(days=1), 0))
<class 'datetime.timedelta'>
>>> type(np.take(datetime.date.today(), 0))
<class 'datetime.date'>

Works for the others though. There might be diminishing returns getting this to work until such a time as we have a templating system in place; then it can generate specific overloads like bool -> np.bool_ etc.

Copy link
Member Author

Choose a reason for hiding this comment

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

If thus is only exception to the rule, then I feel it might be worthwhile to swap generic with an union for the time being.

@overload
def take(
a: _ArrayLike,
indices: int,
axis: Optional[int] = ...,
out: Optional[ndarray] = ...,
mode: _Mode = ...,
) -> generic: ...
@overload
def take(
a: _ArrayLike,
indices: _ArrayLikeInt,
Copy link
Member

Choose a reason for hiding this comment

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

I think _ArrayLikeInt will not end up being general enough because you can do things like

>>> x = np.arange(9).reshape(3, 3)
>>> np.take(x, [[0, 1], [1, 2]], axis=1)
array([[[0, 1],
        [1, 2]],

       [[3, 4],
        [4, 5]],

       [[6, 7],
        [7, 8]]])

Copy link
Member Author

Choose a reason for hiding this comment

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

That's true, I forgot the Sequence[int] syntax does not support arbitrary levels of nesting (which is what we need here).
I don't have a real solution for this, surprisingly fundamental, problem besides adding Unions up to some arbitrary level of nesting that we can consider "good enough".

Copy link
Member

Choose a reason for hiding this comment

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

So far we've just been doing Sequence (i.e. Sequence[Any]). I think the long-term fix is MyPy et. al. supporting recursive types; until then we're stuck with being imprecise.

It's not the best, but I think it's what we have right now. (Though I have also considered the

adding Unions up to some arbitrary level of nesting

It felt too gross to do, but maybe there's more discussion to be had on it in an issue.)

Copy link
Member Author

Choose a reason for hiding this comment

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

It felt too gross to do, but maybe there's more discussion to be had on it in an issue.)

Gross it most definetly is; I've changed it to Sequence[Any] for now (837eec6).

axis: Optional[int] = ...,
out: Optional[ndarray] = ...,
mode: _Mode = ...,
) -> Union[generic, ndarray]: ...
def reshape(a: _ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ...
@overload
def choose(
a: _ScalarGeneric,
choices: Union[Sequence[_ArrayLike], ndarray], # TODO: ndarray[_ArrayLike]
out: Optional[ndarray] = ...,
mode: _Mode = ...,
) -> _ScalarGeneric: ...
@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: ...
62 changes: 62 additions & 0 deletions tests/fail/fromnumeric.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
"""Tests for :mod:`numpy.core.fromnumeric`."""

import numpy as np
Copy link
Member

Choose a reason for hiding this comment

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

Maybe there's a better name than fromnumeric.py? Maybe not though. At some point we're going to need to organize things better, but it's not hard to shuffle stuff around later.

Copy link
Member Author

Choose a reason for hiding this comment

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

Currently it's named after the corresponding module in NumPy (np.core.fromnumeric).
I'm all ears if you have suggestions, but I feel that the current name is already pretty decent.

Copy link
Member Author

Choose a reason for hiding this comment

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

Im any case, I've added a module-level docstring to the tests for now to clarify that the relevant functions are from np.core.fromnumeric (37cbded).


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
65 changes: 65 additions & 0 deletions tests/pass/fromnumeric.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)
67 changes: 67 additions & 0 deletions tests/reveal/fromnumeric.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
"""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: 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