TYP: partition shape-typing (#31169)

jorenham · web-flow · commit f2cffb4ee77b · 2026-04-06T20:34:02.000-06:00
diff --git a/numpy/_core/fromnumeric.pyi b/numpy/_core/fromnumeric.pyi
@@ -345,30 +345,46 @@ def matrix_transpose[ScalarT: np.generic](x: _ArrayLike[ScalarT], /) -> NDArray[
 def matrix_transpose(x: ArrayLike, /) -> NDArray[Any]: ...
 
 #
-@overload
+@overload  # Nd
+def partition[ArrayT: np.ndarray](
+    a: ArrayT,
+    kth: _ArrayLikeInt,
+    axis: SupportsIndex = -1,
+    kind: _PartitionKind = "introselect",
+    order: str | Sequence[str] | None = None,
+) -> ArrayT: ...
+@overload  # ?d
 def partition[ScalarT: np.generic](
     a: _ArrayLike[ScalarT],
     kth: _ArrayLikeInt,
-    axis: SupportsIndex | None = -1,
+    axis: SupportsIndex = -1,
     kind: _PartitionKind = "introselect",
-    order: None = None,
+    order: str | Sequence[str] | None = None,
 ) -> NDArray[ScalarT]: ...
-@overload
-def partition(
-    a: _ArrayLike[np.void],
+@overload  # axis: None
+def partition[ScalarT: np.generic](
+    a: _ArrayLike[ScalarT],
     kth: _ArrayLikeInt,
-    axis: SupportsIndex | None = -1,
+    axis: None,
     kind: _PartitionKind = "introselect",
     order: str | Sequence[str] | None = None,
-) -> NDArray[np.void]: ...
-@overload
+) -> _Array1D[ScalarT]: ...
+@overload  # fallback
 def partition(
     a: ArrayLike,
     kth: _ArrayLikeInt,
-    axis: SupportsIndex | None = -1,
+    axis: SupportsIndex = -1,
     kind: _PartitionKind = "introselect",
     order: str | Sequence[str] | None = None,
 ) -> NDArray[Any]: ...
+@overload  # fallback, axis: None
+def partition(
+    a: ArrayLike,
+    kth: _ArrayLikeInt,
+    axis: None,
+    kind: _PartitionKind = "introselect",
+    order: str | Sequence[str] | None = None,
+) -> _Array1D[Any]: ...
 
 # keep roughly in sync with `ndarray.argpartition`
 @overload  # axis: None
diff --git a/numpy/typing/tests/data/reveal/fromnumeric.pyi b/numpy/typing/tests/data/reveal/fromnumeric.pyi
@@ -81,11 +81,13 @@ assert_type(np.transpose(AR_f4), npt.NDArray[np.float32])
 assert_type(np.transpose(AR_f4_1d), np.ndarray[tuple[int], np.dtype[np.float32]])
 assert_type(np.transpose(AR_f4_2d), np.ndarray[tuple[int, int], np.dtype[np.float32]])
 
-assert_type(np.partition(b, 0, axis=None), npt.NDArray[np.bool])
-assert_type(np.partition(f4, 0, axis=None), npt.NDArray[np.float32])
-assert_type(np.partition(f, 0, axis=None), npt.NDArray[Any])
 assert_type(np.partition(AR_b, 0), npt.NDArray[np.bool])
+assert_type(np.partition(AR_b, 0, axis=None), np.ndarray[tuple[int], np.dtype[np.bool]])
 assert_type(np.partition(AR_f4, 0), npt.NDArray[np.float32])
+assert_type(np.partition(AR_f4, 0, axis=None), np.ndarray[tuple[int], np.dtype[np.float32]])
+assert_type(np.partition(AR_f4_1d, 0), np.ndarray[tuple[int], np.dtype[np.float32]])
+assert_type(np.partition(AR_f4_2d, 0), np.ndarray[tuple[int, int], np.dtype[np.float32]])
+assert_type(np.partition(AR_f4_3d, 0), np.ndarray[tuple[int, int, int], np.dtype[np.float32]])
 
 assert_type(np.argpartition(b, 0), npt.NDArray[np.intp])
 assert_type(np.argpartition(f4, 0), npt.NDArray[np.intp])
