TYP: Shape-typed np.resize annotations

Author: jorenham

numpy/_core/fromnumeric.pyi

@@ -105,6 +105,7 @@ __all__ = [
 _SCT = TypeVar("_SCT", bound=generic)
 _SCT_uifcO = TypeVar("_SCT_uifcO", bound=number[Any] | object_)
 _ArrayType = TypeVar("_ArrayType", bound=np.ndarray[Any, Any])
+_SizeType = TypeVar("_SizeType", bound=int)
 _ShapeType = TypeVar("_ShapeType", bound=tuple[int, ...])
 _ShapeType_co = TypeVar("_ShapeType_co", bound=tuple[int, ...], covariant=True)
 
@@ -427,16 +428,23 @@ def searchsorted(
     sorter: None | _ArrayLikeInt_co = ...,  # 1D int array
 ) -> NDArray[intp]: ...
 
+# unlike `reshape`, `resize` only accepts positive integers, so literal ints can be used
 @overload
-def resize(
-    a: _ArrayLike[_SCT],
-    new_shape: _ShapeLike,
-) -> NDArray[_SCT]: ...
+def resize(a: _ArrayLike[_SCT], new_shape: _SizeType) -> np.ndarray[tuple[_SizeType], np.dtype[_SCT]]: ...
 @overload
-def resize(
-    a: ArrayLike,
-    new_shape: _ShapeLike,
-) -> NDArray[Any]: ...
+def resize(a: _ArrayLike[_SCT], new_shape: SupportsIndex) -> np.ndarray[tuple[int], np.dtype[_SCT]]: ...
+@overload
+def resize(a: _ArrayLike[_SCT], new_shape: _ShapeType) -> np.ndarray[_ShapeType, np.dtype[_SCT]]: ...
+@overload
+def resize(a: _ArrayLike[_SCT], new_shape: Sequence[SupportsIndex]) -> NDArray[_SCT]: ...
+@overload
+def resize(a: ArrayLike, new_shape: _SizeType) -> np.ndarray[tuple[_SizeType], np.dtype[Any]]: ...
+@overload
+def resize(a: ArrayLike, new_shape: SupportsIndex) -> np.ndarray[tuple[int], np.dtype[Any]]: ...
+@overload
+def resize(a: ArrayLike, new_shape: _ShapeType) -> np.ndarray[_ShapeType, np.dtype[Any]]: ...
+@overload
+def resize(a: ArrayLike, new_shape: Sequence[SupportsIndex]) -> NDArray[Any]: ...
 
 @overload
 def squeeze(

numpy/typing/tests/data/reveal/fromnumeric.pyi

@@ -1,6 +1,6 @@
 """Tests for :mod:`_core.fromnumeric`."""
 
-from typing import Any, NoReturn
+from typing import Any, Literal as L, NoReturn
 
 import numpy as np
 import numpy.typing as npt
@@ -102,11 +102,11 @@ assert_type(np.searchsorted(AR_f4[0], 0), np.intp)
 assert_type(np.searchsorted(AR_b[0], [0]), npt.NDArray[np.intp])
 assert_type(np.searchsorted(AR_f4[0], [0]), npt.NDArray[np.intp])
 
-assert_type(np.resize(b, (5, 5)), npt.NDArray[np.bool])
-assert_type(np.resize(f4, (5, 5)), npt.NDArray[np.float32])
-assert_type(np.resize(f, (5, 5)), npt.NDArray[Any])
-assert_type(np.resize(AR_b, (5, 5)), npt.NDArray[np.bool])
-assert_type(np.resize(AR_f4, (5, 5)), npt.NDArray[np.float32])
+assert_type(np.resize(b, (5, 5)), np.ndarray[tuple[L[5], L[5]], np.dtype[np.bool]])
+assert_type(np.resize(f4, (5, 5)), np.ndarray[tuple[L[5], L[5]], np.dtype[np.float32]])
+assert_type(np.resize(f, (5, 5)), np.ndarray[tuple[L[5], L[5]], np.dtype[Any]])
+assert_type(np.resize(AR_b, (5, 5)), np.ndarray[tuple[L[5], L[5]], np.dtype[np.bool]])
+assert_type(np.resize(AR_f4, (5, 5)), np.ndarray[tuple[L[5], L[5]], np.dtype[np.float32]])
 
 assert_type(np.squeeze(b), np.bool)
 assert_type(np.squeeze(f4), np.float32)