Skip to content
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
34 changes: 31 additions & 3 deletions numpy/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -2401,7 +2401,17 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
axis1: SupportsIndex = ...,
axis2: SupportsIndex = ...,
dtype: DTypeLike = ...,
out: _ArrayT = ...,
*,
out: _ArrayT,
) -> _ArrayT: ...
@overload
def trace(
self, # >= 2D array
offset: SupportsIndex,
axis1: SupportsIndex,
axis2: SupportsIndex,
dtype: DTypeLike,
out: _ArrayT,
) -> _ArrayT: ...

@overload
Expand All @@ -2425,7 +2435,16 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
self,
indices: _ArrayLikeInt_co,
axis: SupportsIndex | None = ...,
out: _ArrayT = ...,
*,
out: _ArrayT,
mode: _ModeKind = ...,
) -> _ArrayT: ...
@overload
def take(
self,
indices: _ArrayLikeInt_co,
axis: SupportsIndex | None,
out: _ArrayT,
mode: _ModeKind = ...,
) -> _ArrayT: ...

Expand Down Expand Up @@ -3655,7 +3674,16 @@ class generic(_ArrayOrScalarCommon, Generic[_ItemT_co]):
self,
indices: _ArrayLikeInt_co,
axis: SupportsIndex | None = ...,
out: _ArrayT = ...,
*,
out: _ArrayT,
mode: _ModeKind = ...,
) -> _ArrayT: ...
@overload
def take(
self,
indices: _ArrayLikeInt_co,
axis: SupportsIndex | None,
out: _ArrayT,
mode: _ModeKind = ...,
) -> _ArrayT: ...

Expand Down
62 changes: 56 additions & 6 deletions numpy/_core/multiarray.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -528,7 +528,6 @@ def concatenate( # type: ignore[misc]
casting: _CastingKind | None = ...
) -> NDArray[_ScalarT]: ...
@overload
@overload
def concatenate( # type: ignore[misc]
arrays: SupportsLenAndGetItem[ArrayLike],
/,
Expand All @@ -553,7 +552,17 @@ def concatenate(
arrays: SupportsLenAndGetItem[ArrayLike],
/,
axis: SupportsIndex | None = ...,
out: _ArrayT = ...,
*,
out: _ArrayT,
dtype: DTypeLike = ...,
casting: _CastingKind | None = ...
) -> _ArrayT: ...
@overload
def concatenate(
arrays: SupportsLenAndGetItem[ArrayLike],
/,
axis: SupportsIndex | None,
out: _ArrayT,
*,
dtype: DTypeLike = ...,
casting: _CastingKind | None = ...
Expand Down Expand Up @@ -1094,7 +1103,17 @@ def busday_count(
weekmask: ArrayLike = ...,
holidays: ArrayLike | dt.date | _NestedSequence[dt.date] | None = ...,
busdaycal: busdaycalendar | None = ...,
out: _ArrayT = ...,
*,
out: _ArrayT,
) -> _ArrayT: ...
@overload
def busday_count(
begindates: ArrayLike | dt.date | _NestedSequence[dt.date],
enddates: ArrayLike | dt.date | _NestedSequence[dt.date],
weekmask: ArrayLike,
holidays: ArrayLike | dt.date | _NestedSequence[dt.date] | None,
busdaycal: busdaycalendar | None,
out: _ArrayT,
) -> _ArrayT: ...

# `roll="raise"` is (more or less?) equivalent to `casting="safe"`
Expand Down Expand Up @@ -1126,7 +1145,18 @@ def busday_offset( # type: ignore[misc]
weekmask: ArrayLike = ...,
holidays: ArrayLike | dt.date | _NestedSequence[dt.date] | None = ...,
busdaycal: busdaycalendar | None = ...,
out: _ArrayT = ...,
*,
out: _ArrayT,
) -> _ArrayT: ...
@overload
def busday_offset( # type: ignore[misc]
dates: _ArrayLike[datetime64] | dt.date | _NestedSequence[dt.date],
offsets: _ArrayLikeTD64_co | dt.timedelta | _NestedSequence[dt.timedelta],
roll: L["raise"],
weekmask: ArrayLike,
holidays: ArrayLike | dt.date | _NestedSequence[dt.date] | None,
busdaycal: busdaycalendar | None,
out: _ArrayT,
) -> _ArrayT: ...
@overload
def busday_offset( # type: ignore[misc]
Expand Down Expand Up @@ -1156,7 +1186,18 @@ def busday_offset(
weekmask: ArrayLike = ...,
holidays: ArrayLike | dt.date | _NestedSequence[dt.date] | None = ...,
busdaycal: busdaycalendar | None = ...,
out: _ArrayT = ...,
*,
out: _ArrayT,
) -> _ArrayT: ...
@overload
def busday_offset(
dates: ArrayLike | dt.date | _NestedSequence[dt.date],
offsets: ArrayLike | dt.timedelta | _NestedSequence[dt.timedelta],
roll: _RollKind,
weekmask: ArrayLike,
holidays: ArrayLike | dt.date | _NestedSequence[dt.date] | None,
busdaycal: busdaycalendar | None,
out: _ArrayT,
) -> _ArrayT: ...

@overload
Expand All @@ -1181,7 +1222,16 @@ def is_busday(
weekmask: ArrayLike = ...,
holidays: ArrayLike | dt.date | _NestedSequence[dt.date] | None = ...,
busdaycal: busdaycalendar | None = ...,
out: _ArrayT = ...,
*,
out: _ArrayT,
) -> _ArrayT: ...
@overload
def is_busday(
dates: ArrayLike | _NestedSequence[dt.date],
weekmask: ArrayLike,
holidays: ArrayLike | dt.date | _NestedSequence[dt.date] | None,
busdaycal: busdaycalendar | None,
out: _ArrayT,
) -> _ArrayT: ...

@overload
Expand Down
2 changes: 1 addition & 1 deletion numpy/typing/tests/data/reveal/array_constructors.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ assert_type(np.empty([1, 5, 6], dtype='c16'), npt.NDArray[Any])
assert_type(np.empty(mixed_shape), npt.NDArray[np.float64])

assert_type(np.concatenate(A), npt.NDArray[np.float64])
assert_type(np.concatenate([A, A]), Any) # pyright correctly infers this as NDArray[float64]
assert_type(np.concatenate([A, A]), npt.NDArray[Any]) # pyright correctly infers this as NDArray[float64]
assert_type(np.concatenate([[1], A]), npt.NDArray[Any])
assert_type(np.concatenate([[1], [1]]), npt.NDArray[Any])
assert_type(np.concatenate((A, A)), npt.NDArray[np.float64])
Expand Down
Loading