Merge pull request #28758 from charris/backport-28747

charris · web-flow · commit c200e8cf59bd · 2025-04-18T07:05:38.000-06:00
TYP: some ``[arg]partition`` fixes
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -2293,14 +2293,47 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DType_co]):
         where: _ArrayLikeBool_co = True,
     ) -> _ArrayT: ...
 
+    #
+    @overload
+    def partition(
+        self,
+        /,
+        kth: _ArrayLikeInt,
+        axis: SupportsIndex = -1,
+        kind: _PartitionKind = "introselect",
+        order: None = None,
+    ) -> None: ...
+    @overload
+    def partition(
+        self: NDArray[void],
+        /,
+        kth: _ArrayLikeInt,
+        axis: SupportsIndex = -1,
+        kind: _PartitionKind = "introselect",
+        order: str | Sequence[str] | None = None,
+    ) -> None: ...
+
+    #
+    @overload
     def argpartition(
         self,
-        kth: _ArrayLikeInt_co,
-        axis: None | SupportsIndex = ...,
-        kind: _PartitionKind = ...,
-        order: None | str | Sequence[str] = ...,
+        /,
+        kth: _ArrayLikeInt,
+        axis: SupportsIndex | None = -1,
+        kind: _PartitionKind = "introselect",
+        order: None = None,
+    ) -> NDArray[intp]: ...
+    @overload
+    def argpartition(
+        self: NDArray[void],
+        /,
+        kth: _ArrayLikeInt,
+        axis: SupportsIndex | None = -1,
+        kind: _PartitionKind = "introselect",
+        order: str | Sequence[str] | None = None,
     ) -> NDArray[intp]: ...
 
+    #
     def diagonal(
         self,
         offset: SupportsIndex = ...,
@@ -2320,14 +2353,6 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DType_co]):
     # `nonzero()` is deprecated for 0d arrays/generics
     def nonzero(self) -> tuple[NDArray[intp], ...]: ...
 
-    def partition(
-        self,
-        kth: _ArrayLikeInt_co,
-        axis: SupportsIndex = ...,
-        kind: _PartitionKind = ...,
-        order: None | str | Sequence[str] = ...,
-    ) -> None: ...
-
     # `put` is technically available to `generic`,
     # but is pointless as `generic`s are immutable
     def put(self, /, indices: _ArrayLikeInt_co, values: ArrayLike, mode: _ModeKind = "raise") -> None: ...
diff --git a/numpy/_core/fromnumeric.pyi b/numpy/_core/fromnumeric.pyi
@@ -47,11 +47,11 @@ from numpy._typing import (
     _ShapeLike,
     _ArrayLikeBool_co,
     _ArrayLikeUInt_co,
+    _ArrayLikeInt,
     _ArrayLikeInt_co,
     _ArrayLikeFloat_co,
     _ArrayLikeComplex_co,
     _ArrayLikeObject_co,
-    _ArrayLikeTD64_co,
     _IntLike_co,
     _BoolLike_co,
     _ComplexLike_co,
@@ -322,31 +322,42 @@ def matrix_transpose(x: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
 @overload
 def matrix_transpose(x: ArrayLike, /) -> NDArray[Any]: ...
 
+#
 @overload
 def partition(
     a: _ArrayLike[_SCT],
-    kth: _ArrayLikeInt_co,
-    axis: SupportsIndex | None = ...,
-    kind: _PartitionKind = ...,
-    order: str | Sequence[str] | None = ...,
+    kth: _ArrayLikeInt,
+    axis: SupportsIndex | None = -1,
+    kind: _PartitionKind = "introselect",
+    order: None = None,
 ) -> NDArray[_SCT]: ...
 @overload
+def partition(
+    a: _ArrayLike[np.void],
+    kth: _ArrayLikeInt,
+    axis: SupportsIndex | None = -1,
+    kind: _PartitionKind = "introselect",
+    order: str | Sequence[str] | None = None,
+) -> NDArray[np.void]: ...
+@overload
 def partition(
     a: ArrayLike,
-    kth: _ArrayLikeInt_co,
-    axis: SupportsIndex | None = ...,
-    kind: _PartitionKind = ...,
-    order: str | Sequence[str] | None = ...,
+    kth: _ArrayLikeInt,
+    axis: SupportsIndex | None = -1,
+    kind: _PartitionKind = "introselect",
+    order: str | Sequence[str] | None = None,
 ) -> NDArray[Any]: ...
 
+#
 def argpartition(
     a: ArrayLike,
-    kth: _ArrayLikeInt_co,
+    kth: _ArrayLikeInt,
     axis: SupportsIndex | None = -1,
-    kind: _PartitionKind = ...,
-    order: str | Sequence[str] | None = ...,
+    kind: _PartitionKind = "introselect",
+    order: str | Sequence[str] | None = None,
 ) -> NDArray[intp]: ...
 
+#
 @overload
 def sort(
     a: _ArrayLike[_SCT],
