Merge pull request numpy#27177 from jorenham/typing/arange-shape-type

TYP: 1-d ``numpy.arange`` return shape-type

Author: charris

numpy/_core/multiarray.pyi

@@ -6,6 +6,7 @@ from collections.abc import Sequence, Callable, Iterable
 from typing import (
     Literal as L,
     Any,
+    TypeAlias,
     overload,
     TypeVar,
     SupportsIndex,
@@ -88,6 +89,9 @@ _ArrayType_co = TypeVar(
     bound=ndarray[Any, Any],
     covariant=True,
 )
+_SizeType = TypeVar("_SizeType", bound=int)
+
+_1DArray: TypeAlias = ndarray[tuple[_SizeType], dtype[_SCT]]
 
 # Valid time units
 _UnitKind = L[
@@ -769,7 +773,7 @@ def arange(  # type: ignore[misc]
     dtype: None = ...,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[signedinteger[Any]]: ...
+) -> _1DArray[int, signedinteger[Any]]: ...
 @overload
 def arange(  # type: ignore[misc]
     start: _IntLike_co,
@@ -779,15 +783,15 @@ def arange(  # type: ignore[misc]
     *,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[signedinteger[Any]]: ...
+) -> _1DArray[int, signedinteger[Any]]: ...
 @overload
 def arange(  # type: ignore[misc]
     stop: _FloatLike_co,
     /, *,
     dtype: None = ...,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[floating[Any]]: ...
+) -> _1DArray[int, floating[Any]]: ...
 @overload
 def arange(  # type: ignore[misc]
     start: _FloatLike_co,
@@ -797,15 +801,15 @@ def arange(  # type: ignore[misc]
     *,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[floating[Any]]: ...
+) -> _1DArray[int, floating[Any]]: ...
 @overload
 def arange(
     stop: _TD64Like_co,
     /, *,
     dtype: None = ...,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[timedelta64]: ...
+) -> _1DArray[int, timedelta64]: ...
 @overload
 def arange(
     start: _TD64Like_co,
@@ -815,7 +819,7 @@ def arange(
     *,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[timedelta64]: ...
+) -> _1DArray[int, timedelta64]: ...
 @overload
 def arange(  # both start and stop must always be specified for datetime64
     start: datetime64,
@@ -825,15 +829,15 @@ def arange(  # both start and stop must always be specified for datetime64
     *,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[datetime64]: ...
+) -> _1DArray[int, datetime64]: ...
 @overload
 def arange(
     stop: Any,
     /, *,
     dtype: _DTypeLike[_SCT],
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[_SCT]: ...
+) -> _1DArray[int, _SCT]: ...
 @overload
 def arange(
     start: Any,
@@ -843,15 +847,15 @@ def arange(
     *,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[_SCT]: ...
+) -> _1DArray[int, _SCT]: ...
 @overload
 def arange(
     stop: Any, /,
     *,
     dtype: DTypeLike,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[Any]: ...
+) -> _1DArray[int, Any]: ...
 @overload
 def arange(
     start: Any,
@@ -861,7 +865,7 @@ def arange(
     *,
     device: None | L["cpu"] = ...,
     like: None | _SupportsArrayFunc = ...,
-) -> NDArray[Any]: ...
+) -> _1DArray[int, Any]: ...
 
 def datetime_data(
     dtype: str | _DTypeLike[datetime64] | _DTypeLike[timedelta64], /,

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

@@ -1,5 +1,5 @@
 import sys
-from typing import Any, TypeVar
+from typing import Any, Literal as L, TypeVar
 from pathlib import Path
 from collections import deque
 
@@ -105,18 +105,18 @@ assert_type(np.frombuffer(A), npt.NDArray[np.float64])
 assert_type(np.frombuffer(A, dtype=np.int64), npt.NDArray[np.int64])
 assert_type(np.frombuffer(A, dtype="c16"), npt.NDArray[Any])
 
-assert_type(np.arange(False, True), npt.NDArray[np.signedinteger[Any]])
-assert_type(np.arange(10), npt.NDArray[np.signedinteger[Any]])
-assert_type(np.arange(0, 10, step=2), npt.NDArray[np.signedinteger[Any]])
-assert_type(np.arange(10.0), npt.NDArray[np.floating[Any]])
-assert_type(np.arange(start=0, stop=10.0), npt.NDArray[np.floating[Any]])
-assert_type(np.arange(np.timedelta64(0)), npt.NDArray[np.timedelta64])
-assert_type(np.arange(0, np.timedelta64(10)), npt.NDArray[np.timedelta64])
-assert_type(np.arange(np.datetime64("0"), np.datetime64("10")), npt.NDArray[np.datetime64])
-assert_type(np.arange(10, dtype=np.float64), npt.NDArray[np.float64])
-assert_type(np.arange(0, 10, step=2, dtype=np.int16), npt.NDArray[np.int16])
-assert_type(np.arange(10, dtype=int), npt.NDArray[Any])
-assert_type(np.arange(0, 10, dtype="f8"), npt.NDArray[Any])
+assert_type(np.arange(False, True), np.ndarray[tuple[int], np.dtype[np.signedinteger[Any]]])
+assert_type(np.arange(10), np.ndarray[tuple[int], np.dtype[np.signedinteger[Any]]])
+assert_type(np.arange(0, 10, step=2), np.ndarray[tuple[int], np.dtype[np.signedinteger[Any]]])
+assert_type(np.arange(10.0), np.ndarray[tuple[int], np.dtype[np.floating[Any]]])
+assert_type(np.arange(start=0, stop=10.0), np.ndarray[tuple[int], np.dtype[np.floating[Any]]])
+assert_type(np.arange(np.timedelta64(0)), np.ndarray[tuple[int], np.dtype[np.timedelta64]])
+assert_type(np.arange(0, np.timedelta64(10)), np.ndarray[tuple[int], np.dtype[np.timedelta64]])
+assert_type(np.arange(np.datetime64("0"), np.datetime64("10")), np.ndarray[tuple[int], np.dtype[np.datetime64]])
+assert_type(np.arange(10, dtype=np.float64), np.ndarray[tuple[int], np.dtype[np.float64]])
+assert_type(np.arange(0, 10, step=2, dtype=np.int16), np.ndarray[tuple[int], np.dtype[np.int16]])
+assert_type(np.arange(10, dtype=int), np.ndarray[tuple[int], np.dtype[Any]])
+assert_type(np.arange(0, 10, dtype="f8"), np.ndarray[tuple[int], np.dtype[Any]])
 
 assert_type(np.require(A), npt.NDArray[np.float64])
 assert_type(np.require(B), SubClass[np.float64])