Skip to content

✨🐛 stub _core._multiarray_umath, and many related fixes & improvements #122

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.

Sign up for GitHub

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 1 commit into from
Feb 11, 2025
Merged

✨🐛 stub _core._multiarray_umath, and many related fixes & improvements #122

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
451 changes: 108 additions & 343 deletions src/numpy-stubs/__init__.pyi
View file
Open in desktop

Large diffs are not rendered by default.

1,686 changes: 1,686 additions & 0 deletions src/numpy-stubs/_core/_multiarray_umath.pyi
View file
Open in desktop

Large diffs are not rendered by default.

1,353 changes: 176 additions & 1,177 deletions src/numpy-stubs/_core/multiarray.pyi
View file
Open in desktop

Large diffs are not rendered by default.

434 changes: 227 additions & 207 deletions src/numpy-stubs/_core/numeric.pyi
View file
Open in desktop

Large diffs are not rendered by default.

115 changes: 85 additions & 30 deletions src/numpy-stubs/_typing/_ufunc.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,19 +25,20 @@ _2Tuple: TypeAlias = tuple[_T, _T]
_3Tuple: TypeAlias = tuple[_T, _T, _T]
_4Tuple: TypeAlias = tuple[_T, _T, _T, _T]

_2PTuple: TypeAlias = tuple[_T, _T, *tuple[_T, ...]]
_3PTuple: TypeAlias = tuple[_T, _T, _T, *tuple[_T, ...]]
_4PTuple: TypeAlias = tuple[_T, _T, _T, _T, *tuple[_T, ...]]
_2PTuple: TypeAlias = tuple[_T, _T, Unpack[tuple[_T, ...]]]
_3PTuple: TypeAlias = tuple[_T, _T, _T, Unpack[tuple[_T, ...]]]
_4PTuple: TypeAlias = tuple[_T, _T, _T, _T, Unpack[tuple[_T, ...]]]

_NTypesT_co = TypeVar("_NTypesT_co", bound=int, covariant=True)
_IdentityT_co = TypeVar("_IdentityT_co", covariant=True)
_IdentityT_co = TypeVar("_IdentityT_co", default=None, covariant=True)
_NameT_co = TypeVar("_NameT_co", bound=LiteralString, covariant=True)
_SignatureT_co = TypeVar("_SignatureT_co", bound=LiteralString, covariant=True)

_NInT_co = TypeVar("_NInT_co", bound=int, covariant=True)
_NOutT_co = TypeVar("_NOutT_co", bound=int, covariant=True)
_OutT_co = TypeVar("_OutT_co", covariant=True)
_OutT1_co = TypeVar("_OutT1_co", covariant=True)
_OutT2_co = TypeVar("_OutT2_co", covariant=True)
_ArrayT = TypeVar("_ArrayT", bound=npt.NDArray[Any])
_ArrayT1 = TypeVar("_ArrayT1", bound=npt.NDArray[Any])
_ArrayT2 = TypeVar("_ArrayT2", bound=npt.NDArray[Any])

@type_check_only
class _SupportsArrayUFunc(Protocol):
Expand Down Expand Up @@ -92,7 +93,7 @@ class _UFunc_Nin1_Nout1(np.ufunc, Generic[_NameT_co, _NTypesT_co, _IdentityT_co]
self,
x1: _ScalarLike_co,
/,
out: None = ...,
out: None = None,
*,
where: _ArrayLikeBool_co | None = ...,
casting: np._CastingKind = ...,
Expand Down Expand Up @@ -444,7 +445,7 @@ class _GUFunc_Nin2_Nout1(np.ufunc, Generic[_NameT_co, _NTypesT_co, _IdentityT_co
x1: ArrayLike,
x2: ArrayLike,
/,
out: None = ...,
out: None = None,
*,
casting: np._CastingKind = ...,
order: np._OrderKACF = ...,
Expand Down Expand Up @@ -529,15 +530,15 @@ class _PyFunc_Nin1_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
self,
x1: _ScalarLike_co,
/,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs2],
) -> _OutT_co: ...
@overload
def __call__(
self,
x1: ArrayLike,
/,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs2],
) -> _OutT_co | npt.NDArray[np.object_]: ...
@overload
Expand Down Expand Up @@ -582,7 +583,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
x1: _ScalarLike_co,
x2: _ScalarLike_co,
/,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3],
) -> _OutT_co: ...
@overload
Expand All @@ -591,7 +592,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
x1: ArrayLike,
x2: ArrayLike,
/,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3],
) -> _OutT_co | npt.NDArray[np.object_]: ...
@overload
Expand Down Expand Up @@ -654,7 +655,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
array: ArrayLike,
axis: _ShapeLike | None = ...,
dtype: DTypeLike = ...,
out: None = ...,
out: None = None,
*,
keepdims: Literal[True],
initial: _ScalarLike_co = ...,
Expand All @@ -667,7 +668,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
array: ArrayLike,
axis: _ShapeLike | None = ...,
dtype: DTypeLike = ...,
out: None = ...,
out: None = None,
keepdims: bool = ...,
initial: _ScalarLike_co = ...,
where: _ArrayLikeBool_co = ...,
Expand Down Expand Up @@ -701,7 +702,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
indices: _ArrayLikeInt_co,
axis: SupportsIndex = ...,
dtype: DTypeLike = ...,
out: None = ...,
out: None = None,
) -> npt.NDArray[np.object_]: ...
@overload
def reduceat(
Expand Down Expand Up @@ -738,7 +739,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
array: ArrayLike,
axis: SupportsIndex = ...,
dtype: DTypeLike = ...,
out: None = ...,
out: None = None,
) -> npt.NDArray[np.object_]: ...
@overload
def outer(
Expand All @@ -747,7 +748,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
B: _ScalarLike_co,
/,
*,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3],
) -> _OutT_co: ...
@overload
Expand All @@ -757,7 +758,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
B: ArrayLike,
/,
*,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3],
) -> _OutT_co | npt.NDArray[np.object_]: ...
@overload
Expand All @@ -777,7 +778,7 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
B: _SupportsArrayUFunc | ArrayLike,
/,
*,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3],
) -> Any: ...
@overload
Expand All @@ -787,16 +788,16 @@ class _PyFunc_Nin2_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: i
B: _SupportsArrayUFunc | ArrayLike,
/,
*,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3],
) -> Any: ...

@type_check_only
class _PyFunc_Nin3P_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co, _NInT_co]): # type: ignore[misc] # pyright: ignore[reportGeneralTypeIssues]
class _PyFunc_Nin3P_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: ignore[misc] # pyright: ignore[reportGeneralTypeIssues]
@property
def identity(self) -> _IdentityT_co: ...
@property
def nin(self) -> _NInT_co: ...
def nin(self) -> int: ...
@property
def nout(self) -> Literal[1]: ...
@property
Expand All @@ -811,7 +812,7 @@ class _PyFunc_Nin3P_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co, _NInT_co]):
x3: _ScalarLike_co,
/,
*xs: _ScalarLike_co,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs4P],
) -> _OutT_co: ...
@overload
Expand All @@ -822,7 +823,7 @@ class _PyFunc_Nin3P_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co, _NInT_co]):
x3: ArrayLike,
/,
*xs: ArrayLike,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs4P],
) -> _OutT_co | npt.NDArray[np.object_]: ...
@overload
Expand Down Expand Up @@ -854,13 +855,13 @@ class _PyFunc_Nin3P_Nout1(np.ufunc, Generic[_OutT_co, _IdentityT_co, _NInT_co]):
def outer(self, /, *args: Any, **kwargs: Any) -> NoReturn: ...

@type_check_only
class _PyFunc_Nin1P_Nout2P(np.ufunc, Generic[_OutT_co, _IdentityT_co, _NInT_co, _NOutT_co]): # type: ignore[misc] # pyright: ignore[reportGeneralTypeIssues]
class _PyFunc_Nin1P_Nout2(np.ufunc, Generic[_OutT1_co, _OutT2_co, _IdentityT_co]): # type: ignore[misc] # pyright: ignore[reportGeneralTypeIssues]
@property
def identity(self) -> _IdentityT_co: ...
@property
def nin(self) -> _NInT_co: ...
def nin(self) -> int: ...
@property
def nout(self) -> _NOutT_co: ...
def nout(self) -> Literal[2]: ...
@property
def ntypes(self) -> Literal[1]: ...
@property
Expand All @@ -871,7 +872,61 @@ class _PyFunc_Nin1P_Nout2P(np.ufunc, Generic[_OutT_co, _IdentityT_co, _NInT_co,
x1: _ScalarLike_co,
/,
*xs: _ScalarLike_co,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3P],
) -> tuple[_OutT1_co, _OutT2_co]: ...
@overload
def __call__(
self,
x1: ArrayLike,
/,
*xs: ArrayLike,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3P],
) -> tuple[_OutT1_co | npt.NDArray[np.object_], _OutT2_co | npt.NDArray[np.object_]]: ...
@overload
def __call__(
self,
x1: ArrayLike,
/,
*xs: ArrayLike,
out: tuple[_ArrayT1, _ArrayT2],
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3P],
) -> tuple[_ArrayT1, _ArrayT2]: ...
@overload
def __call__(
self,
x1: _SupportsArrayUFunc | ArrayLike,
/,
*xs: _SupportsArrayUFunc | ArrayLike,
out: _2PTuple[npt.NDArray[Any]] | None = ...,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3P],
) -> Any: ...
def at(self, /, *args: Any, **kwargs: Any) -> NoReturn: ...
def reduce(self, /, *args: Any, **kwargs: Any) -> NoReturn: ...
def accumulate(self, /, *args: Any, **kwargs: Any) -> NoReturn: ...
def reduceat(self, /, *args: Any, **kwargs: Any) -> NoReturn: ...
def outer(self, /, *args: Any, **kwargs: Any) -> NoReturn: ...

@type_check_only
class _PyFunc_Nin1P_Nout2P(np.ufunc, Generic[_OutT_co, _IdentityT_co]): # type: ignore[misc] # pyright: ignore[reportGeneralTypeIssues]
@property
def identity(self) -> _IdentityT_co: ...
@property
def nin(self) -> int: ...
@property
def nout(self) -> int: ...
@property
def ntypes(self) -> Literal[1]: ...
@property
def signature(self) -> None: ...
@overload
def __call__(
self,
x1: _ScalarLike_co,
/,
*xs: _ScalarLike_co,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3P],
) -> _2PTuple[_OutT_co]: ...
@overload
Expand All @@ -880,7 +935,7 @@ class _PyFunc_Nin1P_Nout2P(np.ufunc, Generic[_OutT_co, _IdentityT_co, _NInT_co,
x1: ArrayLike,
/,
*xs: ArrayLike,
out: None = ...,
out: None = None,
**kwargs: Unpack[_PyFunc_Kwargs_Nargs3P],
) -> _2PTuple[_OutT_co | npt.NDArray[np.object_]]: ...
@overload
Expand Down
7 changes: 3 additions & 4 deletions test/runtime/accept/literal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,9 @@
(KACF, AR.flatten),
(KACF, AR.ravel),
(KACF, partial(np.array, 1)),
# NOTE: __call__ is needed due to mypy>=1.11 bugs (#17620, #17631) (also in 1.14)
(CF, partial(np.zeros.__call__, 1)),
(CF, partial(np.ones.__call__, 1)),
(CF, partial(np.empty.__call__, 1)),
(CF, partial(np.zeros, 1)),
(CF, partial(np.ones, 1)),
(CF, partial(np.empty, 1)),
(CF, partial(np.full, 1, 1)),
(KACF, partial(np.zeros_like, AR)),
(KACF, partial(np.ones_like, AR)),
Expand Down
2 changes: 1 addition & 1 deletion test/runtime/accept/recfunctions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ def test_repack_fields() -> None:
assert_type(rfn.repack_fields(dt), np.dtype[np.void])
assert_type(rfn.repack_fields(dt.type(0)), np.void)
assert_type(
rfn.repack_fields(np.ones((int(3),), dtype=dt)),
rfn.repack_fields(np.ones((3,), dtype=dt)),
np.ndarray[tuple[int], np.dtype[np.void]],
)

Expand Down
26 changes: 13 additions & 13 deletions test/static/accept/array_constructors.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import numpy.typing as npt

_SCT_co = TypeVar("_SCT_co", bound=np.generic, covariant=True)

class MyArray(npt.NDArray[_SCT_co]): ...
class MyArray(np.ndarray[tuple[int], np.dtype[_SCT_co]]): ...

i8: np.int64

Expand All @@ -27,7 +27,8 @@ assert_type(np.empty_like(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.empty_like(A, dtype="c16"), npt.NDArray[Any])

assert_type(np.array(A), npt.NDArray[np.float64])
assert_type(np.array(B), npt.NDArray[np.float64])
assert_type(np.array(B), np.ndarray[tuple[int], np.dtype[np.float64]])
assert_type(np.array(D), np.ndarray[tuple[int], np.dtype[np.float64 | np.int64]])
assert_type(np.array([1, 1.0]), npt.NDArray[Any])
assert_type(np.array(deque([1, 2, 3])), npt.NDArray[Any])
assert_type(np.array(A, dtype=np.int64), npt.NDArray[np.int64])
Expand All @@ -37,7 +38,6 @@ assert_type(np.array(A, subok=True), npt.NDArray[np.float64])
assert_type(np.array(B, subok=True), MyArray[np.float64])
assert_type(np.array(B, subok=True, ndmin=0), MyArray[np.float64])
assert_type(np.array(B, subok=True, ndmin=1), MyArray[np.float64])
assert_type(np.array(D), npt.NDArray[np.float64 | np.int64])

assert_type(np.zeros([1, 5, 6]), npt.NDArray[np.float64])
assert_type(np.zeros([1, 5, 6], dtype=np.int64), npt.NDArray[np.int64])
Expand All @@ -59,7 +59,7 @@ assert_type(np.concatenate(A, dtype="c16"), npt.NDArray[Any])
assert_type(np.concatenate([1, 1.0], out=A), npt.NDArray[np.float64])

assert_type(np.asarray(A), npt.NDArray[np.float64])
assert_type(np.asarray(B), npt.NDArray[np.float64])
assert_type(np.asarray(B), np.ndarray[tuple[int], np.dtype[np.float64]])
assert_type(np.asarray([1, 1.0]), npt.NDArray[Any])
assert_type(np.asarray(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.asarray(A, dtype="c16"), npt.NDArray[Any])
Expand All @@ -71,13 +71,13 @@ assert_type(np.asanyarray(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.asanyarray(A, dtype="c16"), npt.NDArray[Any])

assert_type(np.ascontiguousarray(A), npt.NDArray[np.float64])
assert_type(np.ascontiguousarray(B), npt.NDArray[np.float64])
assert_type(np.ascontiguousarray(B), np.ndarray[tuple[int], np.dtype[np.float64]])
assert_type(np.ascontiguousarray([1, 1.0]), npt.NDArray[Any])
assert_type(np.ascontiguousarray(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.ascontiguousarray(A, dtype="c16"), npt.NDArray[Any])

assert_type(np.asfortranarray(A), npt.NDArray[np.float64])
assert_type(np.asfortranarray(B), npt.NDArray[np.float64])
assert_type(np.asfortranarray(B), np.ndarray[tuple[int], np.dtype[np.float64]])
assert_type(np.asfortranarray([1, 1.0]), npt.NDArray[Any])
assert_type(np.asfortranarray(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.asfortranarray(A, dtype="c16"), npt.NDArray[Any])
Expand All @@ -104,17 +104,17 @@ assert_type(np.frombuffer(A), npt.NDArray[np.float64])
assert_type(np.frombuffer(A, dtype=np.int64), npt.NDArray[np.int64])
assert_type(np.frombuffer(A, dtype="c16"), npt.NDArray[Any])

assert_type(np.arange(False, True), np.ndarray[tuple[int], np.dtype[np.signedinteger]])
assert_type(np.arange(10), np.ndarray[tuple[int], np.dtype[np.signedinteger]])
assert_type(np.arange(0, 10, step=2), np.ndarray[tuple[int], np.dtype[np.signedinteger]])
assert_type(np.arange(10.0), np.ndarray[tuple[int], np.dtype[np.floating]])
assert_type(np.arange(start=0, stop=10.0), np.ndarray[tuple[int], np.dtype[np.floating]])
assert_type(np.arange(False, True), np.ndarray[tuple[int], np.dtype[np.int_]])
assert_type(np.arange(10), np.ndarray[tuple[int], np.dtype[np.int_]])
assert_type(np.arange(0, 10, step=2), np.ndarray[tuple[int], np.dtype[np.int_]])
assert_type(np.arange(10.0), np.ndarray[tuple[int], np.dtype[np.float64 | np.int_]])
assert_type(np.arange(start=0, stop=10.0), np.ndarray[tuple[int], np.dtype[np.float64 | np.int_]])
assert_type(np.arange(np.timedelta64(0)), np.ndarray[tuple[int], np.dtype[np.timedelta64]])
assert_type(np.arange(0, np.timedelta64(10)), np.ndarray[tuple[int], np.dtype[np.timedelta64]])
assert_type(np.arange(np.datetime64("0"), np.datetime64("10")), np.ndarray[tuple[int], np.dtype[np.datetime64]])
assert_type(np.arange(10, dtype=np.float64), np.ndarray[tuple[int], np.dtype[np.float64]])
assert_type(np.arange(0, 10, step=2, dtype=np.int16), np.ndarray[tuple[int], np.dtype[np.int16]])
assert_type(np.arange(10, dtype=int), np.ndarray[tuple[int], np.dtype[Any]])
assert_type(np.arange(10, dtype=int), np.ndarray[tuple[int], np.dtype[np.int_]])
assert_type(np.arange(0, 10, dtype="f8"), np.ndarray[tuple[int], np.dtype[Any]])

assert_type(np.require(A), npt.NDArray[np.float64])
Expand Down Expand Up @@ -179,7 +179,7 @@ assert_type(np.ones(_shape_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]
assert_type(np.ones(_shape_nd), np.ndarray[tuple[int, ...], np.dtype[np.float64]])
assert_type(np.ones(_shape_1d, dtype=np.int64), np.ndarray[tuple[int], np.dtype[np.int64]])
assert_type(np.ones(_shape_like), npt.NDArray[np.float64])
assert_type(np.ones(_shape_like, dtype=np.dtypes.Int64DType()), np.ndarray[Any, np.dtypes.Int64DType])
assert_type(np.ones(_shape_like, dtype=np.dtypes.Int64DType()), npt.NDArray[np.int64])
assert_type(np.ones(_shape_like, dtype=int), npt.NDArray[Any])

assert_type(np.full(_size, i8), np.ndarray[tuple[int], np.dtype[np.int64]])
Expand Down
Loading
Loading