TYP: fix overloads where out: _ArrayT was typed as being the default (#29278)

MarcoGorelli · web-flow · commit 6eeaae25ddfa · 2025-06-26T20:42:09.000+02:00
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -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
@@ -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: ...
 
@@ -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: ...
 
diff --git a/numpy/_core/multiarray.pyi b/numpy/_core/multiarray.pyi
@@ -528,7 +528,6 @@ def concatenate(  # type: ignore[misc]
     casting: _CastingKind | None = ...
 ) -> NDArray[_ScalarT]: ...
 @overload
-@overload
 def concatenate(  # type: ignore[misc]
     arrays: SupportsLenAndGetItem[ArrayLike],
     /,
@@ -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 = ...
@@ -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"`
@@ -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]
@@ -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
@@ -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
diff --git a/numpy/typing/tests/data/reveal/array_constructors.pyi b/numpy/typing/tests/data/reveal/array_constructors.pyi
@@ -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])
