Skip to content

MAINT: Update the annotations in np.core.numeric #17564

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Oct 21, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
120 changes: 96 additions & 24 deletions numpy/core/numeric.pyi
Original file line number Diff line number Diff line change
@@ -1,14 +1,45 @@
from typing import Any, Optional, Union, Sequence, Tuple
import sys
from typing import (
Any,
Optional,
Union,
Sequence,
Tuple,
Callable,
List,
overload,
TypeVar,
Iterable,
)

from numpy import ndarray, dtype, bool_, _OrderKACF, _OrderCF
from numpy import ndarray, generic, dtype, bool_, int32, int64, _OrderKACF, _OrderCF
from numpy.typing import ArrayLike, DtypeLike, _ShapeLike

if sys.version_info >= (3, 8):
from typing import Literal
else:
from typing_extensions import Literal

_T = TypeVar("_T")
_ArrayType = TypeVar("_ArrayType", bound=ndarray)

_CorrelateMode = Literal["valid", "same", "full"]

@overload
def zeros_like(
a: _ArrayType,
dtype: None = ...,
order: _OrderKACF = ...,
subok: Literal[True] = ...,
shape: None = ...,
) -> _ArrayType: ...
@overload
def zeros_like(
a: ArrayLike,
dtype: DtypeLike = ...,
order: _OrderKACF = ...,
subok: bool = ...,
shape: Optional[Union[int, Sequence[int]]] = ...,
shape: Optional[_ShapeLike] = ...,
) -> ndarray: ...
def ones(
shape: _ShapeLike,
Expand All @@ -17,13 +48,31 @@ def ones(
*,
like: ArrayLike = ...,
) -> ndarray: ...
@overload
def ones_like(
a: _ArrayType,
dtype: None = ...,
order: _OrderKACF = ...,
subok: Literal[True] = ...,
shape: None = ...,
) -> _ArrayType: ...
@overload
def ones_like(
a: ArrayLike,
dtype: DtypeLike = ...,
order: _OrderKACF = ...,
subok: bool = ...,
shape: Optional[_ShapeLike] = ...,
) -> ndarray: ...
@overload
def empty_like(
a: _ArrayType,
dtype: None = ...,
order: _OrderKACF = ...,
subok: Literal[True] = ...,
shape: None = ...,
) -> _ArrayType: ...
@overload
def empty_like(
a: ArrayLike,
dtype: DtypeLike = ...,
Expand All @@ -39,6 +88,16 @@ def full(
*,
like: ArrayLike = ...,
) -> ndarray: ...
@overload
def full_like(
a: _ArrayType,
fill_value: Any,
dtype: None = ...,
order: _OrderKACF = ...,
subok: Literal[True] = ...,
shape: None = ...,
) -> _ArrayType: ...
@overload
def full_like(
a: ArrayLike,
fill_value: Any,
Expand All @@ -47,35 +106,38 @@ def full_like(
subok: bool = ...,
shape: Optional[_ShapeLike] = ...,
) -> ndarray: ...
@overload
def count_nonzero(
a: ArrayLike, axis: None = ..., *, keepdims: Literal[False] = ...
) -> int: ...
Copy link
Member Author

Choose a reason for hiding this comment

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

A bit odd that one overload returns a builtin int and the other one a numpy signedinteger (or an array).
See #17562.

@overload
def count_nonzero(
a: ArrayLike, axis: Optional[Union[int, Tuple[int], Tuple[int, int]]] = ...
) -> Union[int, ndarray]: ...
def isfortran(a: ndarray) -> bool: ...
a: ArrayLike, axis: _ShapeLike = ..., *, keepdims: bool = ...
) -> Union[int64, int32, ndarray]: ...
def isfortran(a: Union[ndarray, generic]) -> bool: ...
def argwhere(a: ArrayLike) -> ndarray: ...
def flatnonzero(a: ArrayLike) -> ndarray: ...

_CorrelateMode = Literal["valid", "same", "full"]

def correlate(a: ArrayLike, v: ArrayLike, mode: _CorrelateMode = ...) -> ndarray: ...
def convolve(a: ArrayLike, v: ArrayLike, mode: _CorrelateMode = ...) -> ndarray: ...
def outer(a: ArrayLike, b: ArrayLike, out: ndarray = ...) -> ndarray: ...
@overload
def outer(a: ArrayLike, b: ArrayLike, out: None = ...) -> ndarray: ...
@overload
def outer(a: ArrayLike, b: ArrayLike, out: _ArrayType = ...) -> _ArrayType: ...
def tensordot(
a: ArrayLike,
b: ArrayLike,
axes: Union[
int, Tuple[int, int], Tuple[Tuple[int, int], ...], Tuple[List[int, int], ...]
] = ...,
axes: Union[int, Tuple[_ShapeLike, _ShapeLike]] = ...,
) -> ndarray: ...
def roll(
a: ArrayLike,
shift: Union[int, Tuple[int, ...]],
axis: Optional[Union[int, Tuple[int, ...]]] = ...,
shift: _ShapeLike,
axis: Optional[_ShapeLike] = ...,
) -> ndarray: ...
def rollaxis(a: ArrayLike, axis: int, start: int = ...) -> ndarray: ...
def rollaxis(a: ndarray, axis: int, start: int = ...) -> ndarray: ...
Copy link
Member Author

Choose a reason for hiding this comment

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

As mentioned in the docs: rollaxis takes an array, not an array-like object.

def moveaxis(
a: ndarray,
source: Union[int, Sequence[int]],
destination: Union[int, Sequence[int]],
source: _ShapeLike,
destination: _ShapeLike,
) -> ndarray: ...
def cross(
a: ArrayLike,
Expand All @@ -85,16 +147,26 @@ def cross(
axisc: int = ...,
axis: Optional[int] = ...,
) -> ndarray: ...
@overload
def indices(
dimensions: Sequence[int], dtype: dtype = ..., sparse: bool = ...
) -> Union[ndarray, Tuple[ndarray, ...]]: ...
dimensions: Sequence[int],
dtype: DtypeLike = ...,
sparse: Literal[False] = ...,
) -> ndarray: ...
@overload
def indices(
dimensions: Sequence[int],
dtype: DtypeLike = ...,
sparse: Literal[True] = ...,
) -> Tuple[ndarray, ...]: ...
def fromfunction(
function: Callable,
shape: Tuple[int, int],
function: Callable[..., _T],
shape: Sequence[int],
*,
dtype: DtypeLike = ...,
like: ArrayLike = ...,
**kwargs,
) -> Any: ...
**kwargs: Any,
) -> _T: ...
def isscalar(element: Any) -> bool: ...
def binary_repr(num: int, width: Optional[int] = ...) -> str: ...
def base_repr(number: int, base: int = ..., padding: int = ...) -> str: ...
Expand Down
40 changes: 39 additions & 1 deletion numpy/typing/tests/data/pass/array_constructors.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import List
from typing import List, Any
import numpy as np

class Index:
Expand All @@ -7,10 +7,15 @@ def __index__(self) -> int:

class SubClass(np.ndarray): ...

i8 = np.int64(1)

A = np.array([1])
B = A.view(SubClass).copy()
C = [1]

def func(i: int, j: int, **kwargs: Any) -> SubClass:
return B

np.array(1, dtype=float)
np.array(1, copy=False)
np.array(1, order='F')
Expand Down Expand Up @@ -64,3 +69,36 @@ class SubClass(np.ndarray): ...
np.logspace(0, 2, base=[1j, 2j], num=2)

np.geomspace(1, 2)

np.zeros_like(A)
np.zeros_like(C)
np.zeros_like(B)
np.zeros_like(B, dtype=np.int64)

np.ones_like(A)
np.ones_like(C)
np.ones_like(B)
np.ones_like(B, dtype=np.int64)

np.empty_like(A)
np.empty_like(C)
np.empty_like(B)
np.empty_like(B, dtype=np.int64)

np.full_like(A, i8)
np.full_like(C, i8)
np.full_like(B, i8)
np.full_like(B, i8, dtype=np.int64)

np.ones(1)
np.ones([1, 1, 1])

np.full(1, i8)
np.full([1, 1, 1], i8)

np.indices([1, 2, 3])
np.indices([1, 2, 3], sparse=True)

np.fromfunction(func, (3, 5))

np.identity(10)
89 changes: 89 additions & 0 deletions numpy/typing/tests/data/pass/numeric.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
"""
Tests for :mod:`numpy.core.numeric`.

Does not include tests which fall under ``array_constructors``.

"""

from typing import List
import numpy as np

class SubClass(np.ndarray):
...

i8 = np.int64(1)

A = np.arange(27).reshape(3, 3, 3)
B: List[List[List[int]]] = A.tolist()
C = np.empty((27, 27)).view(SubClass)

np.count_nonzero(i8)
np.count_nonzero(A)
np.count_nonzero(B)
np.count_nonzero(A, keepdims=True)
np.count_nonzero(A, axis=0)

np.isfortran(i8)
np.isfortran(A)

np.argwhere(i8)
np.argwhere(A)

np.flatnonzero(i8)
np.flatnonzero(A)

np.correlate(B[0][0], A.ravel(), mode="valid")
np.correlate(A.ravel(), A.ravel(), mode="same")

np.convolve(B[0][0], A.ravel(), mode="valid")
np.convolve(A.ravel(), A.ravel(), mode="same")

np.outer(i8, A)
np.outer(B, A)
np.outer(A, A)
np.outer(A, A, out=C)

np.tensordot(B, A)
np.tensordot(A, A)
np.tensordot(A, A, axes=0)
np.tensordot(A, A, axes=(0, 1))

np.isscalar(i8)
np.isscalar(A)
np.isscalar(B)

np.roll(A, 1)
np.roll(A, (1, 2))
np.roll(B, 1)

np.rollaxis(A, 0, 1)

np.moveaxis(A, 0, 1)
np.moveaxis(A, (0, 1), (1, 2))

np.cross(B, A)
np.cross(A, A)

np.indices([0, 1, 2])
np.indices([0, 1, 2], sparse=False)
np.indices([0, 1, 2], sparse=True)

np.binary_repr(1)

np.base_repr(1)

np.allclose(i8, A)
np.allclose(B, A)
np.allclose(A, A)

np.isclose(i8, A)
np.isclose(B, A)
np.isclose(A, A)

np.array_equal(i8, A)
np.array_equal(B, A)
np.array_equal(A, A)

np.array_equiv(i8, A)
np.array_equiv(B, A)
np.array_equiv(A, A)
8 changes: 0 additions & 8 deletions numpy/typing/tests/data/pass/simple.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,6 @@ def ndarray_func(x):
array == 1
array.dtype == float

ndarray_func(np.zeros([1, 2]))
ndarray_func(np.ones([1, 2]))
ndarray_func(np.empty([1, 2]))

ndarray_func(np.zeros_like(array))
ndarray_func(np.ones_like(array))
ndarray_func(np.empty_like(array))

# Dtype construction
np.dtype(float)
np.dtype(np.float64)
Expand Down
Loading