From 6376466ebcf039a6328f579f9bfaba68db06222d Mon Sep 17 00:00:00 2001 From: jorenham Date: Thu, 2 Apr 2026 23:59:58 +0200 Subject: [PATCH] TYP: ``transpose`` and ``matrix_transpose`` shape-typing --- numpy/_core/fromnumeric.pyi | 17 ++++++++--------- numpy/ma/core.pyi | 6 ++++-- numpy/typing/tests/data/reveal/fromnumeric.pyi | 4 ++++ 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/numpy/_core/fromnumeric.pyi b/numpy/_core/fromnumeric.pyi index 132c6ae73468..27f48f7fe5c3 100644 --- a/numpy/_core/fromnumeric.pyi +++ b/numpy/_core/fromnumeric.pyi @@ -311,18 +311,17 @@ def swapaxes[ScalarT: np.generic](a: _ArrayLike[ScalarT], axis1: SupportsIndex, @overload def swapaxes(a: ArrayLike, axis1: SupportsIndex, axis2: SupportsIndex) -> NDArray[Any]: ... -# keep in sync with `ma.core.transpose` +# @overload -def transpose[ScalarT: np.generic]( - a: _ArrayLike[ScalarT], - axes: _ShapeLike | None = None, -) -> NDArray[ScalarT]: ... +def transpose[ArrayT: np.ndarray](a: ArrayT, axes: _ShapeLike | None = None) -> ArrayT: ... @overload -def transpose( - a: ArrayLike, - axes: _ShapeLike | None = None, -) -> NDArray[Any]: ... +def transpose[ScalarT: np.generic](a: _ArrayLike[ScalarT], axes: _ShapeLike | None = None) -> NDArray[ScalarT]: ... +@overload +def transpose(a: ArrayLike, axes: _ShapeLike | None = None) -> NDArray[Any]: ... +# +@overload +def matrix_transpose[ArrayT: np.ndarray](x: ArrayT, /) -> ArrayT: ... @overload def matrix_transpose[ScalarT: np.generic](x: _ArrayLike[ScalarT], /) -> NDArray[ScalarT]: ... @overload diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi index d7b92e12065b..13df91ac4c36 100644 --- a/numpy/ma/core.pyi +++ b/numpy/ma/core.pyi @@ -3571,9 +3571,11 @@ def putmask(a: np.ndarray, mask: _ArrayLikeBool_co, values: ArrayLike) -> None: # keep in sync with `_core.fromnumeric.transpose` @overload -def transpose[ScalarT: np.generic](a: _ArrayLike[ScalarT], axes: _ShapeLike | None = None) -> _MaskedArray[ScalarT]: ... +def transpose[ArrayT: np.ndarray](a: ArrayT, axes: _ShapeLike | None = None) -> ArrayT: ... @overload -def transpose(a: ArrayLike, axes: _ShapeLike | None = None) -> _MaskedArray[Incomplete]: ... +def transpose[ScalarT: np.generic](a: _ArrayLike[ScalarT], axes: _ShapeLike | None = None) -> _MaskedArray[ScalarT]: ... +@overload # `_MaskedArray | np.ndarray` is equivalent to `np.ndarray` +def transpose(a: ArrayLike, axes: _ShapeLike | None = None) -> np.ndarray: ... # keep in sync with `_core.fromnumeric.reshape` @overload # shape: index diff --git a/numpy/typing/tests/data/reveal/fromnumeric.pyi b/numpy/typing/tests/data/reveal/fromnumeric.pyi index 5afaf6cc84a6..89060b456349 100644 --- a/numpy/typing/tests/data/reveal/fromnumeric.pyi +++ b/numpy/typing/tests/data/reveal/fromnumeric.pyi @@ -9,6 +9,8 @@ class NDArraySubclass(np.ndarray[tuple[Any, ...], np.dtype[np.complex128]]): ... AR_b: npt.NDArray[np.bool] AR_f4: npt.NDArray[np.float32] +AR_f4_1d: np.ndarray[tuple[int], np.dtype[np.float32]] +AR_f4_2d: np.ndarray[tuple[int, int], np.dtype[np.float32]] AR_c16: npt.NDArray[np.complex128] AR_u8: npt.NDArray[np.uint64] AR_i8: npt.NDArray[np.int64] @@ -68,6 +70,8 @@ assert_type(np.transpose(f4), npt.NDArray[np.float32]) assert_type(np.transpose(f), npt.NDArray[Any]) assert_type(np.transpose(AR_b), npt.NDArray[np.bool]) 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])